const Texture = @This();
const std = @import("std");

const vk = @import("vulkan");

const CommandBuffer = @import("CommandBuffer.zig").CommandBuffer;
const Engine = @import("Engine.zig");
const StagingBuffer = @import("StagingBuffer.zig");
const TargetQueue = @import("TargetQueue.zig").TargetQueue;

pub const Usage = enum {
    base_color,
    normal,
    occlusion_roughness_metallic,
    emissive,
    depth,

    pub fn vkFormat(self: Usage) vk.Format {
        return switch (self) {
            .base_color => .r8g8b8a8_srgb,
            .normal => .r8g8b8a8_snorm,
            .occlusion_roughness_metallic => .r8g8b8a8_unorm,
            .emissive => .r16g16b16a16_sfloat,
            .depth => .d32_sfloat,
        };
    }

    pub fn samplesPerTexel(self: Usage) u32 {
        return switch (self) {
            .base_color => 4,
            .normal => 4,
            .occlusion_roughness_metallic => 4,
            .emissive => 4,
            .depth => 1,
        };
    }

    pub fn SampleType(comptime self: Usage) type {
        return switch (self) {
            .base_color => u8,
            .normal => i8,
            .occlusion_roughness_metallic => u8,
            .emissive => f16,
            .depth => f32,
        };
    }

    pub fn bytesPerSample(self: Usage) u32 {
        return switch (self) {
            inline else => |x| @sizeOf(SampleType(x)),
        };
    }

    pub fn TexelType(comptime self: Usage) type {
        return [self.samplesPerTexel()]SampleType(self);
    }

    pub fn bytesPerTexel(self: Usage) u32 {
        return switch (self) {
            inline else => |x| @sizeOf(TexelType(x)),
        };
    }
};

pub const InitInfo = struct {
    width: u32,
    height: u32,
    usage: Usage,
    target_queue: TargetQueue,
    name: ?[]const u8 = null,
};

image: vk.Image,
image_view: vk.ImageView,
device_memory: vk.DeviceMemory,
target_queue: TargetQueue,

width: u32,
height: u32,
usage: Usage,

pub fn init(engine: *Engine, init_info: InitInfo) !Texture {
    const target_queue_family = switch (init_info.target_queue) {
        .graphics => engine.graphics_queue.allocation.family,
        .compute => engine.compute_queue.allocation.family,
    };
    const transfer_queue_family = engine.transfer_queue.allocation.family;

    const queue_family_indices: []const u32 = if (init_info.usage == .depth)
        &.{target_queue_family}
    else
        &.{ target_queue_family, transfer_queue_family };

    const image = try engine.createImage(.{
        .image_type = .@"2d",
        .format = init_info.usage.vkFormat(),
        .extent = .{
            .width = init_info.width,
            .height = init_info.height,
            .depth = 1,
        },
        .mip_levels = 1,
        .array_layers = 1,
        .samples = .{ .@"1_bit" = true },
        .tiling = .optimal,
        .usage = .{
            .depth_stencil_attachment_bit = init_info.usage == .depth,
            .transfer_dst_bit = init_info.usage != .depth,
            .sampled_bit = init_info.usage != .depth,
        },
        .queue_family_indices = queue_family_indices,
        .initial_layout = .undefined,
    });
    errdefer engine.destroyImage(image);
    if (init_info.name) |name| {
        engine.setObjectName(image, "I {s} [{s}]", .{ name, @tagName(init_info.usage) });
    }

    const memory_requirements = engine.device.getImageMemoryRequirements(image);
    const device_memory = try engine.allocate(memory_requirements, .{ .device_local_bit = true });
    errdefer engine.freeMemory(device_memory);
    if (init_info.name) |name| {
        engine.setObjectName(device_memory, "DM {s} [{s}]", .{ name, @tagName(init_info.usage) });
    }

    try engine.device.bindImageMemory(image, device_memory, 0);

    const image_view = try engine.createImageView(.{
        .image = image,
        .view_type = .@"2d",
        .format = init_info.usage.vkFormat(),
        .subresource_range = .{
            .aspect_mask = .{
                .color_bit = init_info.usage != .depth,
                .depth_bit = init_info.usage == .depth,
            },
            .base_mip_level = 0,
            .level_count = 1,
            .base_array_layer = 0,
            .layer_count = 1,
        },
    });
    errdefer engine.destroyImageView(image_view);
    if (init_info.name) |name| {
        engine.setObjectName(image_view, "IV {s} [{s}]", .{ name, @tagName(init_info.usage) });
    }

    return .{
        .image = image,
        .image_view = image_view,
        .device_memory = device_memory,
        .target_queue = init_info.target_queue,
        .width = init_info.width,
        .height = init_info.height,
        .usage = init_info.usage,
    };
}

pub fn deinit(self: *Texture, engine: *Engine) void {
    std.log.debug("Deinitializing {*} with {*}", .{ self, engine });

    engine.destroyImageView(self.image_view);
    engine.freeMemory(self.device_memory);
    engine.destroyImage(self.image);

    self.* = undefined;
}

pub fn writeTexels(self: Texture, comptime TexelType: type, engine: *Engine, texels: []const TexelType) !void {
    const texel_count = try std.math.mul(u32, self.width, self.height);
    std.debug.assert(texels.len == texel_count);
    switch (self.usage) {
        inline else => |x| std.debug.assert(TexelType == x.TexelType()),
    }

    try self.writeRaw(engine, std.mem.sliceAsBytes(texels));
}

pub fn writeSamples(self: Texture, comptime SampleType: type, engine: *Engine, samples: []const SampleType) !void {
    const texel_count = try std.math.mul(u32, self.width, self.height);
    const sample_count = try std.math.mul(u32, texel_count, self.usage.samplesPerTexel());
    std.debug.assert(samples.len == sample_count);
    switch (self.usage) {
        inline else => |x| std.debug.assert(SampleType == x.SampleType()),
    }

    try self.writeRaw(engine, std.mem.sliceAsBytes(samples));
}

pub fn writeRaw(self: Texture, engine: *Engine, data: []const u8) !void {
    const texel_count = try std.math.mul(u32, self.width, self.height);
    const byte_length = try std.math.mul(u32, texel_count, self.usage.bytesPerTexel());
    std.debug.assert(data.len == byte_length);
    std.debug.assert(self.usage != .depth);

    var staging_buffer = try StagingBuffer.init(engine, .{
        .capacity = @intCast(byte_length),
        .target_queue = self.target_queue,
    });
    defer staging_buffer.deinit(engine);

    const staging_memory = try staging_buffer.map(engine);
    @memcpy(staging_memory, data);
    staging_buffer.unmap(engine);

    // --- TRANSITION TO TRANSFER_DST_OPTIMAL AND COPY -----------------

    var transfer_command_buffer: CommandBuffer(.transfer, .transient) = try .init(engine);
    defer transfer_command_buffer.deinit(engine);

    try transfer_command_buffer.beginCommandBuffer(.{});

    transfer_command_buffer.pipelineBarrier(.{
        .src_stage_mask = .{ .top_of_pipe_bit = true },
        .dst_stage_mask = .{ .transfer_bit = true },
        .image_memory_barriers = &.{
            .{
                .src_access_mask = .{},
                .dst_access_mask = .{ .transfer_write_bit = true },
                .old_layout = .undefined,
                .new_layout = .transfer_dst_optimal,
                .src_queue_family_index = vk.QUEUE_FAMILY_IGNORED,
                .dst_queue_family_index = vk.QUEUE_FAMILY_IGNORED,
                .image = self.image,
                .subresource_range = .{
                    .aspect_mask = .{ .color_bit = true },
                    .base_mip_level = 0,
                    .level_count = 1,
                    .base_array_layer = 0,
                    .layer_count = 1,
                },
            },
        },
    });

    transfer_command_buffer.copyBufferToImage(
        staging_buffer.buffer,
        self.image,
        .transfer_dst_optimal,
        &.{
            .{
                .buffer_offset = 0,
                .buffer_row_length = self.width,
                .buffer_image_height = self.height,
                .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 = self.width, .height = self.height, .depth = 1 },
            },
        },
    );

    try transfer_command_buffer.endCommandBuffer();

    const semaphore = try engine.createSemaphore();
    defer engine.destroySemaphore(semaphore);

    try transfer_command_buffer.submit(engine, .{
        .signal_semaphores = &.{semaphore},
    });

    // --- TRANSITION TO SHADER_READ_ONLY_OPTIMAL ----------------------

    var graphics_command_buffer: CommandBuffer(.graphics, .transient) = try .init(engine);
    defer graphics_command_buffer.deinit(engine);

    try graphics_command_buffer.beginCommandBuffer(.{});

    graphics_command_buffer.pipelineBarrier(.{
        .src_stage_mask = .{ .transfer_bit = true },
        .dst_stage_mask = .{ .fragment_shader_bit = true },
        .image_memory_barriers = &.{
            .{
                .src_access_mask = .{ .transfer_write_bit = true },
                .dst_access_mask = .{ .shader_read_bit = true },
                .old_layout = .transfer_dst_optimal,
                .new_layout = .shader_read_only_optimal,
                .src_queue_family_index = vk.QUEUE_FAMILY_IGNORED,
                .dst_queue_family_index = vk.QUEUE_FAMILY_IGNORED,
                .image = self.image,
                .subresource_range = .{
                    .aspect_mask = .{ .color_bit = true },
                    .base_mip_level = 0,
                    .level_count = 1,
                    .base_array_layer = 0,
                    .layer_count = 1,
                },
            },
        },
    });

    try graphics_command_buffer.endCommandBuffer();

    const fence = try engine.createFence(.{});
    defer engine.destroyFence(fence);

    try graphics_command_buffer.submit(engine, .{
        .wait_semaphores = &.{.{ .semaphore = semaphore }},
        .fence = fence,
    });
    try engine.waitForFence(fence);
}