const std = @import("std");

const vk = @import("vulkan");

const Engine = @import("Engine.zig");
const QueueType = @import("QueueType.zig").QueueType;
const Transient = @import("Transient.zig").Transient;
const WaitSemaphore = @import("WaitSemaphore.zig");

pub const DrawIndexed = struct {
    index_count: u32,
    instance_count: u32 = 1,
    first_index: u32 = 0,
    vertex_offset: i32 = 0,
    first_instance: u32 = 0,
};

pub const PipelineBarrier = struct {
    src_stage_mask: vk.PipelineStageFlags,
    dst_stage_mask: vk.PipelineStageFlags,
    dependency_flags: vk.DependencyFlags = .{},
    memory_barriers: []const vk.MemoryBarrier = &.{},
    buffer_memory_barriers: []const vk.BufferMemoryBarrier = &.{},
    image_memory_barriers: []const vk.ImageMemoryBarrier = &.{},
};

pub const RenderPassBeginInfo = struct {
    render_pass: vk.RenderPass,
    framebuffer: vk.Framebuffer,
    render_area: vk.Rect2D,
    clear_values: []const vk.ClearValue = &.{},
};

pub const SubmitInfo = struct {
    wait_semaphores: []const WaitSemaphore = &.{},
    signal_semaphores: []const vk.Semaphore = &.{},
    fence: vk.Fence = .null_handle,
};

pub const VertexBufferBinding = struct {
    buffer: vk.Buffer,
    offset: usize,
};

pub fn CommandBuffer(comptime queue_type: QueueType, comptime transient: Transient) type {
    return struct {
        const Self = @This();

        proxy: vk.CommandBufferProxy,
        allocator: std.mem.Allocator,

        pub fn init(engine: *const Engine) !Self {
            const handle = try allocateCommandBuffer(engine, queue_type, transient);
            const proxy: vk.CommandBufferProxy = .init(handle, engine.device.wrapper);
            const allocator = engine.vk_allocator.allocator;
            return .{ .proxy = proxy, .allocator = allocator };
        }

        pub fn deinit(self: *Self, engine: *const Engine) void {
            freeCommandBuffer(engine, queue_type, transient, self.proxy.handle);
            self.* = undefined;
        }

        pub fn submit(self: Self, engine: *const Engine, submit_info: SubmitInfo) !void {
            try submitCommandBuffer(engine, queue_type, self.proxy.handle, submit_info);
        }

        pub const CommandBufferBeginInfo = struct {
            flags: vk.CommandBufferUsageFlags = .{
                .one_time_submit_bit = switch (transient) {
                    .persistent => false,
                    .transient => true,
                },
            },
            inheritance_info: ?vk.CommandBufferInheritanceInfo = null,
        };

        pub fn beginCommandBuffer(self: Self, begin_info: CommandBufferBeginInfo) !void {
            try self.proxy.beginCommandBuffer(&.{
                .flags = begin_info.flags,
                .p_inheritance_info = if (begin_info.inheritance_info) |*x| x else null,
            });
        }

        pub inline fn beginRenderPass(self: Self, render_pass_begin: RenderPassBeginInfo, contents: vk.SubpassContents) void {
            self.proxy.beginRenderPass(&.{
                .render_pass = render_pass_begin.render_pass,
                .framebuffer = render_pass_begin.framebuffer,
                .render_area = render_pass_begin.render_area,
                .clear_value_count = @intCast(render_pass_begin.clear_values.len),
                .p_clear_values = render_pass_begin.clear_values.ptr,
            }, contents);
        }

        pub inline fn bindDescriptorSets(
            self: Self,
            pipeline_bind_point: vk.PipelineBindPoint,
            layout: vk.PipelineLayout,
            first_set: u32,
            descriptor_sets: []const vk.DescriptorSet,
            dynamic_offsets: []const u32,
        ) void {
            self.proxy.bindDescriptorSets(
                pipeline_bind_point,
                layout,
                first_set,
                @intCast(descriptor_sets.len),
                descriptor_sets.ptr,
                @intCast(dynamic_offsets.len),
                dynamic_offsets.ptr,
            );
        }

        pub inline fn bindPipeline(self: Self, pipeline_bind_point: vk.PipelineBindPoint, pipeline: vk.Pipeline) void {
            self.proxy.bindPipeline(pipeline_bind_point, pipeline);
        }

        pub inline fn bindIndexBuffer(self: Self, buffer: vk.Buffer, offset: u64, index_type: vk.IndexType) void {
            self.proxy.bindIndexBuffer(buffer, offset, index_type);
        }

        pub inline fn bindVertexBuffers(self: Self, first_binding: u32, bindings: []const VertexBufferBinding) !void {
            var vertex_buffer_bindings = std.MultiArrayList(VertexBufferBinding){};
            defer vertex_buffer_bindings.deinit(self.allocator);
            try vertex_buffer_bindings.ensureTotalCapacity(self.allocator, bindings.len);

            for (bindings) |binding| {
                vertex_buffer_bindings.appendAssumeCapacity(binding);
            }

            self.proxy.bindVertexBuffers(
                first_binding,
                @intCast(vertex_buffer_bindings.len),
                vertex_buffer_bindings.items(.buffer).ptr,
                vertex_buffer_bindings.items(.offset).ptr,
            );
        }

        pub inline fn drawIndexed(self: Self, draw_indexed: DrawIndexed) void {
            self.proxy.drawIndexed(
                draw_indexed.index_count,
                draw_indexed.instance_count,
                draw_indexed.first_index,
                draw_indexed.vertex_offset,
                draw_indexed.first_instance,
            );
        }

        pub inline fn endCommandBuffer(self: Self) !void {
            try self.proxy.endCommandBuffer();
        }

        pub inline fn endRenderPass(self: Self) void {
            self.proxy.endRenderPass();
        }

        pub inline fn copyBuffer(
            self: Self,
            src_buffer: vk.Buffer,
            dst_buffer: vk.Buffer,
            regions: []const vk.BufferCopy,
        ) void {
            self.proxy.copyBuffer(src_buffer, dst_buffer, @intCast(regions.len), regions.ptr);
        }

        pub inline fn copyBufferToImage(
            self: Self,
            src_buffer: vk.Buffer,
            dst_image: vk.Image,
            dst_image_layout: vk.ImageLayout,
            regions: []const vk.BufferImageCopy,
        ) void {
            self.proxy.copyBufferToImage(src_buffer, dst_image, dst_image_layout, @intCast(regions.len), regions.ptr);
        }

        pub inline fn pipelineBarrier(self: Self, barrier: PipelineBarrier) void {
            self.proxy.pipelineBarrier(
                barrier.src_stage_mask,
                barrier.dst_stage_mask,
                barrier.dependency_flags,
                @intCast(barrier.memory_barriers.len),
                barrier.memory_barriers.ptr,
                @intCast(barrier.buffer_memory_barriers.len),
                barrier.buffer_memory_barriers.ptr,
                @intCast(barrier.image_memory_barriers.len),
                barrier.image_memory_barriers.ptr,
            );
        }

        pub inline fn setScissor(self: Self, first_scissor: u32, scissors: []const vk.Rect2D) void {
            self.proxy.setScissor(first_scissor, @intCast(scissors.len), scissors.ptr);
        }

        pub inline fn setViewport(self: Self, first_viewport: u32, viewports: []const vk.Viewport) void {
            self.proxy.setViewport(first_viewport, @intCast(viewports.len), viewports.ptr);
        }
    };
}

inline fn resolveCommandPool(self: *const Engine, comptime queue_type: QueueType, comptime transient: Transient) vk.CommandPool {
    return switch (queue_type) {
        .graphics => switch (transient) {
            .persistent => self.graphics_command_pool,
            .transient => self.transient_graphics_command_pool,
        },
        .compute => switch (transient) {
            .persistent => self.compute_command_pool,
            .transient => self.transient_compute_command_pool,
        },
        .transfer => switch (transient) {
            .persistent => self.transfer_command_pool,
            .transient => self.transient_transfer_command_pool,
        },
    };
}

fn allocateCommandBuffer(engine: *const Engine, comptime queue_type: QueueType, comptime transient: Transient) !vk.CommandBuffer {
    const command_pool = resolveCommandPool(engine, queue_type, transient);

    var command_buffers: [1]vk.CommandBuffer = undefined;
    try engine.device.allocateCommandBuffers(&.{
        .command_pool = command_pool,
        .level = .primary,
        .command_buffer_count = command_buffers.len,
    }, &command_buffers);

    return command_buffers[0];
}

fn freeCommandBuffer(engine: *const Engine, comptime queue_type: QueueType, comptime transient: Transient, command_buffer: vk.CommandBuffer) void {
    const command_pool = resolveCommandPool(engine, queue_type, transient);
    const command_buffers = [_]vk.CommandBuffer{command_buffer};

    engine.device.freeCommandBuffers(command_pool, command_buffers.len, &command_buffers);
}

fn submitCommandBuffer(
    engine: *const Engine,
    comptime queue_type: QueueType,
    command_buffer: vk.CommandBuffer,
    submit_info: SubmitInfo,
) !void {
    const queue = switch (queue_type) {
        .graphics => engine.graphics_queue.handle,
        .compute => engine.compute_queue.handle,
        .transfer => engine.transfer_queue.handle,
    };

    const command_buffers = [_]vk.CommandBuffer{command_buffer};

    var wait_semaphores = std.MultiArrayList(WaitSemaphore){};
    defer wait_semaphores.deinit(engine.vk_allocator.allocator);
    try wait_semaphores.ensureTotalCapacity(engine.vk_allocator.allocator, submit_info.wait_semaphores.len);

    for (submit_info.wait_semaphores) |wait_semaphore| {
        wait_semaphores.appendAssumeCapacity(wait_semaphore);
    }

    const submits = [_]vk.SubmitInfo{
        .{
            .wait_semaphore_count = @intCast(wait_semaphores.len),
            .p_wait_semaphores = wait_semaphores.items(.semaphore).ptr,
            .p_wait_dst_stage_mask = wait_semaphores.items(.stage_flags).ptr,
            .command_buffer_count = command_buffers.len,
            .p_command_buffers = &command_buffers,
            .signal_semaphore_count = @intCast(submit_info.signal_semaphores.len),
            .p_signal_semaphores = submit_info.signal_semaphores.ptr,
        },
    };

    try engine.device.queueSubmit(queue, submits.len, &submits, submit_info.fence);
}