const vk = @import("vulkan");

const context = @import("context.zig");
const device = @import("device.zig");
const buffer = @import("buffer.zig");
const commands = @import("commands.zig");
const assets = @import("../assets.zig");

pub const TextureContext = struct {
    image: vk.Image,
    memory: vk.DeviceMemory,
    image_view: vk.ImageView,
    sampler: vk.Sampler,

    pub fn destroy(self: *const TextureContext, ldc: *const device.LogicalDeviceContext) void {
        ldc.vkd.destroySampler(ldc.device, self.sampler, null);
        ldc.vkd.destroyImageView(ldc.device, self.image_view, null);
        ldc.vkd.destroyImage(ldc.device, self.image, null);
        ldc.vkd.freeMemory(ldc.device, self.memory, null);
    }
};

pub fn createImage(
    vc: context.VulkanContext,
    ldc: device.LogicalDeviceContext,
    width: u32,
    height: u32,
    format: vk.Format,
    tiling: vk.ImageTiling,
    usage: vk.ImageUsageFlags,
    properties: vk.MemoryPropertyFlags,
) !TextureContext {
    const image_create_info = vk.ImageCreateInfo{
        .image_type = .@"2d",
        .extent = .{
            .width = width,
            .height = height,
            .depth = 1,
        },
        .mip_levels = 1,
        .array_layers = 1,
        .format = format,
        .tiling = tiling,
        .initial_layout = .undefined,
        .usage = usage,
        .samples = .{ .@"1_bit" = true },
        .sharing_mode = .exclusive,
    };

    const image = try ldc.vkd.createImage(ldc.device, &image_create_info, null);
    errdefer ldc.vkd.destroyImage(ldc.device, image, null);

    const memory_requirements = ldc.vkd.getImageMemoryRequirements(ldc.device, image);
    const memory_type_index = try buffer.findMemoryType(
        vc,
        ldc,
        memory_requirements.memory_type_bits,
        properties,
    );

    const alloc_info = vk.MemoryAllocateInfo{
        .allocation_size = memory_requirements.size,
        .memory_type_index = memory_type_index,
    };

    const memory = try ldc.vkd.allocateMemory(ldc.device, &alloc_info, null);
    errdefer ldc.vkd.freeMemory(ldc.device, memory, null);

    try ldc.vkd.bindImageMemory(ldc.device, image, memory, 0);

    const image_view = try createImageView(ldc, image, format);
    errdefer ldc.vkd.destroyImageView(ldc.device, image_view, null);

    const sampler = try createTextureSampler(ldc);
    errdefer ldc.vkd.destroySampler(ldc.device, sampler, null);

    return .{
        .image = image,
        .memory = memory,
        .image_view = image_view,
        .sampler = sampler,
    };
}

fn createImageView(
    ldc: device.LogicalDeviceContext,
    image: vk.Image,
    format: vk.Format,
) !vk.ImageView {
    const view_create_info = vk.ImageViewCreateInfo{
        .image = image,
        .view_type = .@"2d",
        .format = format,
        .components = .{
            .r = .identity,
            .g = .identity,
            .b = .identity,
            .a = .identity,
        },
        .subresource_range = .{
            .aspect_mask = .{
                .color_bit = true,
            },
            .base_mip_level = 0,
            .level_count = 1,
            .base_array_layer = 0,
            .layer_count = 1,
        },
    };

    return try ldc.vkd.createImageView(
        ldc.device,
        &view_create_info,
        null,
    );
}

fn createTextureSampler(ldc: device.LogicalDeviceContext) !vk.Sampler {
    const sampler_create_info = vk.SamplerCreateInfo{
        .mag_filter = .linear,
        .min_filter = .linear,
        .mipmap_mode = .linear,
        .address_mode_u = .clamp_to_edge,
        .address_mode_v = .clamp_to_edge,
        .address_mode_w = .clamp_to_edge,
        .mip_lod_bias = 0.0,
        .anisotropy_enable = .false,
        .max_anisotropy = 1.0,
        .compare_enable = .false,
        .compare_op = .always,
        .min_lod = 0.0,
        .max_lod = 0.0,
        .border_color = .int_opaque_black,
        .unnormalized_coordinates = .false,
    };

    return try ldc.vkd.createSampler(
        ldc.device,
        &sampler_create_info,
        null,
    );
}

pub fn transitionImageLayout(
    ldc: device.LogicalDeviceContext,
    command_context: commands.CommandContext,
    image: vk.Image,
    old_layout: vk.ImageLayout,
    new_layout: vk.ImageLayout,
) !void {
    const command_buffer = try commands.beginSingleTimeCommands(ldc, command_context);

    var barrier = vk.ImageMemoryBarrier{
        .src_access_mask = .{},
        .dst_access_mask = .{},
        .old_layout = old_layout,
        .new_layout = new_layout,
        .src_queue_family_index = vk.QUEUE_FAMILY_IGNORED,
        .dst_queue_family_index = vk.QUEUE_FAMILY_IGNORED,
        .image = image,
        .subresource_range = .{
            .aspect_mask = .{ .color_bit = true },
            .base_mip_level = 0,
            .level_count = 1,
            .base_array_layer = 0,
            .layer_count = 1,
        },
    };

    var source_stage: vk.PipelineStageFlags = undefined;
    var destination_stage: vk.PipelineStageFlags = undefined;

    if (old_layout == .undefined and new_layout == .transfer_dst_optimal) {
        barrier.src_access_mask = .{};
        barrier.dst_access_mask = .{ .transfer_write_bit = true };

        source_stage = .{ .top_of_pipe_bit = true };
        destination_stage = .{ .transfer_bit = true };
    } else if (old_layout == .transfer_dst_optimal and new_layout == .shader_read_only_optimal) {
        barrier.src_access_mask = .{ .transfer_write_bit = true };
        barrier.dst_access_mask = .{ .shader_read_bit = true };

        source_stage = .{ .transfer_bit = true };
        destination_stage = .{ .fragment_shader_bit = true };
    } else {
        return error.UnsupportedLayoutTransition;
    }

    const image_barriers = [_]vk.ImageMemoryBarrier{barrier};

    ldc.vkd.cmdPipelineBarrier(
        command_buffer,
        source_stage,
        destination_stage,
        .{},
        null,
        null,
        &image_barriers,
    );

    try commands.endSingleTimeCommands(ldc, command_context, command_buffer);
}

pub fn copyBufferToImage(
    ldc: device.LogicalDeviceContext,
    command_context: commands.CommandContext,
    source_buffer: vk.Buffer,
    destination_image: vk.Image,
    width: u32,
    height: u32,
) !void {
    const command_buffer = try commands.beginSingleTimeCommands(ldc, command_context);

    const region = vk.BufferImageCopy{
        .buffer_offset = 0,
        .buffer_row_length = 0,
        .buffer_image_height = 0,
        .image_subresource = .{
            .aspect_mask = .{ .color_bit = true },
            .mip_level = 0,
            .base_array_layer = 0,
            .layer_count = 1,
        },
        .image_offset = .{ .x = 0, .y = 0, .z = 0 },
        .image_extent = .{
            .width = width,
            .height = height,
            .depth = 1,
        },
    };

    const regions = [_]vk.BufferImageCopy{region};

    ldc.vkd.cmdCopyBufferToImage(
        command_buffer,
        source_buffer,
        destination_image,
        .transfer_dst_optimal,
        &regions,
    );

    try commands.endSingleTimeCommands(ldc, command_context, command_buffer);
}

pub fn initTextureFromRgba(
    vc: context.VulkanContext,
    ldc: device.LogicalDeviceContext,
    command_context: commands.CommandContext,
    asset: []const u8,
    width: u32,
    height: u32,
    format: vk.Format,
    tiling: vk.ImageTiling,
    usage: vk.ImageUsageFlags,
    properties: vk.MemoryPropertyFlags,
) !TextureContext {
    const staging = try buffer.initStagingBuffer(vc, ldc, asset);
    defer staging.destroy(&ldc);

    const texture = try createImage(
        vc,
        ldc,
        width,
        height,
        format,
        tiling,
        usage,
        properties,
    );
    try transitionImageLayout(
        ldc,
        command_context,
        texture.image,
        .undefined,
        .transfer_dst_optimal,
    );
    try copyBufferToImage(
        ldc,
        command_context,
        staging.buffer,
        texture.image,
        width,
        height,
    );
    try transitionImageLayout(
        ldc,
        command_context,
        texture.image,
        .transfer_dst_optimal,
        .shader_read_only_optimal,
    );
    return texture;
}