const std = @import("std");
const assert = @import("std").debug.assert;

const vk = @import("vulkan");

const settings = @import("settings.zig");
const Instance = @import("instance.zig").Instance;
const Device = @import("device.zig").Device;
const Window = @import("display.zig").Window;
const Renderer = @import("Renderer.zig");
const Command = @import("Command.zig");
const sync = @import("sync.zig");

// TODO memory
var gpa = std.heap.GeneralPurposeAllocator(.{}){};
const allocator = gpa.allocator();

const present_stages: vk.PipelineStageFlags = .{ .all_commands_bit = true, .top_of_pipe_bit = true, .bottom_of_pipe_bit = true };
const present_access: vk.AccessFlags = .{
    .memory_read_bit = true,
    .memory_write_bit = true,
};

const graphics_stages: vk.PipelineStageFlags = .{
    .all_commands_bit = true,
    .all_graphics_bit = true,
    .top_of_pipe_bit = true,
    .bottom_of_pipe_bit = true,
    .draw_indirect_bit = true,
    .vertex_input_bit = true,
    .vertex_shader_bit = true,
    .tessellation_control_shader_bit = true,
    .tessellation_evaluation_shader_bit = true,
    .geometry_shader_bit = true,
    .fragment_shader_bit = true,
    .early_fragment_tests_bit = true,
    .late_fragment_tests_bit = true,
    .color_attachment_output_bit = true,
    .conditional_rendering_bit_ext = true,
    .task_shader_bit_nv = true,
    .mesh_shader_bit_nv = true,
};
const graphics_access: vk.AccessFlags = .{
    .memory_read_bit = true,
    .memory_write_bit = true,
    .indirect_command_read_bit = true,
    .index_read_bit = true,
    .vertex_attribute_read_bit = true,
    .uniform_read_bit = true,
    .input_attachment_read_bit = true,
    .shader_read_bit = true,
    .shader_write_bit = true,
    .color_attachment_read_bit = true,
    .color_attachment_write_bit = true,
    .conditional_rendering_read_bit_ext = true,
    .acceleration_structure_read_bit_khr = true,
    .acceleration_structure_write_bit_khr = true,
};

const compute_stages: vk.PipelineStageFlags = .{
    .all_commands_bit = true,
    .top_of_pipe_bit = true,
    .bottom_of_pipe_bit = true,
    .compute_shader_bit = true,
    .acceleration_structure_build_bit_khr = true,
    .ray_tracing_shader_bit_khr = true,
};
const compute_access: vk.AccessFlags = .{
    .memory_read_bit = true,
    .memory_write_bit = true,
    .indirect_command_read_bit = true,
    .uniform_read_bit = true,
    .shader_read_bit = true,
    .shader_write_bit = true,
    .acceleration_structure_read_bit_khr = true,
    .acceleration_structure_write_bit_khr = true,
};

const transfer_stages: vk.PipelineStageFlags = .{
    .all_commands_bit = true,
    .top_of_pipe_bit = true,
    .bottom_of_pipe_bit = true,
    .transfer_bit = true,
    .host_bit = true,
};
const transfer_access: vk.AccessFlags = .{
    .memory_read_bit = true,
    .memory_write_bit = true,
    .host_read_bit = true,
    .host_write_bit = true,
    .transfer_read_bit = true,
    .transfer_write_bit = true,
};

pub const QueueId = enum(u4) {
    graphics,
    compute,
    transfer,
    present,

    pub const count = @typeInfo(@This()).Enum.fields.len;
};

pub const QueueFlags = packed struct {
    graphics_bit: bool = false,
    compute_bit: bool = false,
    transfer_bit: bool = false,
    present_bit: bool = false,

    __reserved_bits_04_31: u28 = 0,

    pub usingnamespace vk.FlagsMixin(@This(), vk.Flags);
};

var s_queues: [QueueId.count]Queue = undefined;

pub fn init(instance: *const Instance, device: *const Device, window: *const Window) !void {
    const support = try Support.init(instance, device.physical_device, window);
    defer support.deinit();

    for (s_queues) |*queue, i| {
        queue.* = try Queue.init(device, &support, i);
    }
}

pub fn deinit(device: *const Device) void {
    device.waitIdle() catch {};

    for (s_queues) |queue| {
        queue.deinit(device);
    }
}

pub fn get(id: QueueId) *Queue {
    return &s_queues[@enumToInt(id)];
}

pub const Support = struct {
    const Self = @This();

    families: [QueueId.count]?u32 = undefined,
    indices: [QueueId.count]QueueId = undefined,
    properties: []vk.QueueFamilyProperties = undefined,

    pub fn init(instance: *const Instance, physical_device: vk.PhysicalDevice, window: *const Window) !Self {
        var self = Self{};

        var count: u32 = undefined;
        instance.dispatch.getPhysicalDeviceQueueFamilyProperties(physical_device, &count, null);
        self.properties = try allocator.alloc(vk.QueueFamilyProperties, count);
        instance.dispatch.getPhysicalDeviceQueueFamilyProperties(physical_device, &count, self.properties.ptr);

        self.families[@enumToInt(QueueId.graphics)] = selectGfxFamily(self.properties);
        self.families[@enumToInt(QueueId.compute)] = selectCompFamily(self.properties);
        self.families[@enumToInt(QueueId.transfer)] = selectXferFamily(self.properties);
        self.families[@enumToInt(QueueId.present)] = try selectPresFamily(instance, physical_device, window.surface, self.properties);

        // TODO: don't understand the purpose of 'indices' yet...
        const choice_counts = try allocator.alloc(u32, count);
        defer allocator.free(choice_counts);
        for (self.families) |family, i| {
            if (family) |f| {
                const choice_count = choice_counts[@intCast(usize, f)];
                self.indices[i] = @intToEnum(QueueId, choice_count);
                // probably wrong
                choice_counts[@intCast(usize, f)] = choice_count + 1;
            }
        }

        return self;
    }

    pub fn deinit(self: *const Self) void {
        allocator.free(self.properties);
    }

    fn selectGfxFamily(families: []vk.QueueFamilyProperties) ?u32 {
        var index: ?u32 = null;
        var score: i32 = 0;
        for (families) |family, i| {
            if (family.queue_count == 0) {
                continue;
            }

            if (family.queue_flags.graphics_bit) {
                var new_score: i32 = 0;
                new_score += if (family.queue_flags.compute_bit) @as(i32, 1) else 0;
                new_score += if (family.queue_flags.transfer_bit) @as(i32, 1) else 0;
                if (new_score > score) {
                    score = new_score;
                    index = @intCast(u32, i);
                }
            }
        }

        return index;
    }

    fn selectCompFamily(families: []vk.QueueFamilyProperties) ?u32 {
        var index: ?u32 = null;
        var score: i32 = 0;
        for (families) |family, i| {
            if (family.queue_count == 0) {
                continue;
            }

            if (family.queue_flags.compute_bit) {
                var new_score: i32 = 0;
                new_score += if (family.queue_flags.graphics_bit) @as(i32, 1) else 0;
                new_score += if (family.queue_flags.transfer_bit) @as(i32, 1) else 0;
                if (new_score > score) {
                    score = new_score;
                    index = @intCast(u32, i);
                }
            }
        }

        return index;
    }

    fn selectXferFamily(families: []vk.QueueFamilyProperties) ?u32 {
        var index: ?u32 = null;
        var score: i32 = 0;
        for (families) |family, i| {
            if (family.queue_count == 0) {
                continue;
            }

            if (family.queue_flags.transfer_bit) {
                var new_score: i32 = 0;
                new_score += if (family.queue_flags.graphics_bit) @as(i32, 1) else 0;
                new_score += if (family.queue_flags.compute_bit) @as(i32, 1) else 0;
                if (new_score > score) {
                    score = new_score;
                    index = @intCast(u32, i);
                }
            }
        }

        return index;
    }

    fn selectPresFamily(instance: *const Instance, phdev: vk.PhysicalDevice, surf: vk.SurfaceKHR, families: []vk.QueueFamilyProperties) !?u32 {
        var index: ?u32 = null;
        var score: i32 = 0;
        for (families) |family, i| {
            const result = try instance.dispatch.getPhysicalDeviceSurfaceSupportKHR(phdev, @intCast(u32, i), surf);
            const presentable = result != 0;
            if (presentable) {
                var new_score: i32 = 0;
                new_score += if (family.queue_flags.graphics_bit) @as(i32, 1) else 0;
                new_score += if (family.queue_flags.compute_bit) @as(i32, 1) else 0;
                new_score += if (family.queue_flags.transfer_bit) @as(i32, 1) else 0;
                if (new_score > score) {
                    score = new_score;
                    index = @intCast(u32, i);
                }
            }
        }

        return index;
    }
};

pub const Queue = struct {
    handle: vk.Queue = undefined,
    family: u32 = 0,
    index: QueueId = 0,
    access_mask: vk.AccessFlags = .{},
    stage_mask: vk.PipelineStageFlags = .{},
    queueId: u4 = 0,
    gfx: bool = false,
    comp: bool = false,
    xfer: bool = false,
    pres: bool = false,

    cmd_pool: Command.Pool = undefined,
    cmds: [settings.cmds_per_queue]vk.CommandBuffer = undefined,
    cmd_fences: [settings.cmds_per_queue]sync.Fence = undefined,
    cmd_ids: [settings.cmds_per_queue]u32 = undefined,
    head: u32 = 0,
    tail: u32 = 0,

    const Self = @This();
    pub fn init(device: *const Device, support: *const Support, id: usize) !Self {
        const family = support.families[id] orelse unreachable; // fix this silliness
        const index = support.indices[id];
        const handle = device.dispatch.getDeviceQueue(device.handle, family, @enumToInt(index));

        var self = Self{
            .family = family,
            .index = index,
            .handle = handle,
        };

        const pres_family = support.families[@enumToInt(QueueId.present)];
        const queue_flags = support.properties[@enumToInt(QueueId.present)].queue_flags;

        if (queue_flags.graphics_bit) {
            self.gfx = true;
            self.stage_mask = self.stage_mask.merge(graphics_stages);
            self.access_mask = self.access_mask.merge(graphics_access);
        }

        if (queue_flags.compute_bit) {
            self.comp = true;
            self.stage_mask = self.stage_mask.merge(compute_stages);
            self.access_mask = self.access_mask.merge(compute_access);
        }

        if (queue_flags.transfer_bit) {
            self.xfer = true;
            self.stage_mask = self.stage_mask.merge(transfer_stages);
            self.access_mask = self.access_mask.merge(transfer_access);
        }

        if (family == pres_family) {
            self.pres = true;
            self.stage_mask = self.stage_mask.merge(present_stages);
            self.access_mask = self.access_mask.merge(present_access);
        }

        assert(self.stage_mask.toInt() != 0);
        assert(self.access_mask.toInt() != 0);

        try Command.init(device, &self, @intToEnum(QueueId, id));

        return self;
    }

    pub fn deinit(self: *const Self, device: *const Device) void {
        Command.deinit(device, self);
    }
};