const vk = @import("vulkan");
const glfw = @import("glfw");
const log2 = @import("std").math.log2;

const settings = @import("settings.zig");
const OnlyIf = settings.OnlyIf;

const Extensions = @import("Extensions.zig");
const Instance = @import("instance.zig").Instance;
const enabled_layers = @import("layers.zig").enabled;
const Renderer = @import("Renderer.zig");
const Window = @import("display.zig").Window;
const queues = @import("queues.zig");

const std = @import("std");

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

const DeviceDispatch = vk.DeviceWrapper(.{
    .destroyDevice = true,
    .getDeviceQueue = true,
    .createSemaphore = true,
    .createFence = true,
    .createImageView = true,
    .destroyImageView = true,
    .destroySemaphore = true,
    .destroyFence = true,
    .getSwapchainImagesKHR = true,
    .createSwapchainKHR = true,
    .destroySwapchainKHR = true,
    .acquireNextImageKHR = true,
    .deviceWaitIdle = true,
    .waitForFences = true,
    .resetFences = true,
    .queueSubmit = true,
    .queuePresentKHR = true,
    .createCommandPool = true,
    .destroyCommandPool = true,
    .allocateCommandBuffers = true,
    .freeCommandBuffers = true,
    .queueWaitIdle = true,
    .createShaderModule = true,
    .destroyShaderModule = true,
    .createPipelineLayout = true,
    .destroyPipelineLayout = true,
    .createRenderPass = true,
    .destroyRenderPass = true,
    .createGraphicsPipelines = true,
    .destroyPipeline = true,
    .createFramebuffer = true,
    .destroyFramebuffer = true,
    .beginCommandBuffer = true,
    .endCommandBuffer = true,
    .resetCommandBuffer = true,
    .allocateMemory = true,
    .freeMemory = true,
    .createBuffer = true,
    .destroyBuffer = true,
    .getBufferMemoryRequirements = true,
    .mapMemory = true,
    .unmapMemory = true,
    .bindBufferMemory = true,
    .cmdBeginRenderPass = true,
    .cmdEndRenderPass = true,
    .cmdBindPipeline = true,
    .cmdDraw = true,
    .cmdSetViewport = true,
    .cmdSetScissor = true,
    .cmdBindVertexBuffers = true,
    .cmdCopyBuffer = true,
    .getFenceStatus = true,
});

const Props = struct {
    phdev: vk.PhysicalDeviceProperties2,
    accstr: OnlyIf(settings.rt_on, vk.PhysicalDeviceAccelerationStructurePropertiesKHR),
    rtpipe: OnlyIf(settings.rt_on, vk.PhysicalDeviceRayTracingPipelinePropertiesKHR),

    pub fn getName(self: *Props) []u8 {
        return std.mem.sliceTo(&self.phdev.properties.device_name, 0);
    }
};

const Features = struct {
    phdev: vk.PhysicalDeviceFeatures2,
    accstr: OnlyIf(settings.rt_on, vk.PhysicalDeviceAccelerationStructureFeaturesKHR) = undefined,
    rtpipe: OnlyIf(settings.rt_on, vk.PhysicalDeviceRayTracingPipelineFeaturesKHR) = undefined,
    rquery: OnlyIf(settings.rt_on, vk.PhysicalDeviceRayQueryFeaturesKHR) = undefined,
};

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

    physical_device: vk.PhysicalDevice = undefined,
    handle: vk.Device = undefined,
    dispatch: DeviceDispatch = undefined,
    props: Props = undefined,
    exts: Extensions.Device = undefined,
    feats: Features = undefined,

    pub fn init(instance: *const Instance, window: *const Window) !Self {
        var self = Self{};
        try selectPhysicalDevice(&self, instance, window);

        try self.createDevice(window, instance);
        self.dispatch = try DeviceDispatch.load(self.handle, instance.dispatch.dispatch.vkGetDeviceProcAddr);

        errdefer self.dispatch.destroyDevice(self.handle, null);

        // TODO: volk init?

        try queues.init(instance, &self, window);
        return self;
    }

    pub fn deinit(self: *const Self) void {
        queues.deinit(self);
        self.dispatch.destroyDevice(self.handle, null);
    }

    pub fn waitIdle(self: *const Self) !void {
        try self.dispatch.deviceWaitIdle(self.handle);
    }

    pub fn getName(self: *Self) ![]u8 {
        return self.props.getName();
    }

    const DeviceScore = struct {
        rt_score: i32,
        ext_score: i32,
        feat_score: i32,
        prop_score: i32,
        has_required_exts: bool,
        has_queue_support: bool,
    };

    fn selectPhysicalDevice(self: *Device, instance: *const Instance, window: *const Window) !void {
        var device_count: u32 = undefined;
        _ = try instance.dispatch.enumeratePhysicalDevices(instance.handle, &device_count, null);

        const pdevs = try allocator.alloc(vk.PhysicalDevice, device_count);
        defer allocator.free(pdevs);

        _ = try instance.dispatch.enumeratePhysicalDevices(instance.handle, &device_count, pdevs.ptr);

        const prop_list = try allocator.alloc(Props, device_count);
        defer allocator.free(prop_list);

        for (prop_list) |*prop, i| {
            prop.* = .{
                .phdev = .{
                    .s_type = .physical_device_properties_2,
                    .properties = undefined,
                },
                .accstr = undefined,
                .rtpipe = undefined,
            };

            if (settings.rt_on) {
                prop.phdev.p_next = prop.accstr;
                prop.accstr = .{
                    .s_type = .physical_device_acceleration_structure_properties_khr,
                    .p_next = &prop.rtpipe,
                };

                prop.rtpipe = .{
                    .s_type = .physical_device_ray_tracing_pipeline_properties_khr,
                };
            }

            instance.dispatch.getPhysicalDeviceProperties2(pdevs[i], &prop.phdev);
        }

        const feats_list = try allocator.alloc(Features, device_count);
        defer allocator.free(feats_list);

        for (feats_list) |*feat, i| {
            feat.phdev.s_type = .physical_device_features_2;
            if (settings.rt_on) {
                feat.phdev.p_next = &feat.accstr;
                feat.accstr.s_type = .physical_device_acceleration_structure_features_khr;
                feat.accstr.p_next = &feat.rtpipe;
                feat.rtpipe.s_type = .physical_device_ray_tracing_pipeline_features_khr;
                feat.rtpipe.p_next = &feat.rquery;
                feat.rquery.s_type = .physical_device_ray_query_features_khr;
            }

            instance.dispatch.getPhysicalDeviceFeatures2(pdevs[i], &feat.phdev);
        }

        const exts_list = try allocator.alloc(Extensions.Device, device_count);
        defer allocator.free(exts_list);

        for (exts_list) |*ext, i| {
            var count: u32 = undefined;
            _ = try instance.dispatch.enumerateDeviceExtensionProperties(pdevs[i], null, &count, null);
            const props = try allocator.alloc(vk.ExtensionProperties, count);
            defer allocator.free(props);
            _ = try instance.dispatch.enumerateDeviceExtensionProperties(pdevs[i], null, &count, props.ptr);

            ext.* = Extensions.Device.get(props);
        }

        const scores_list = try allocator.alloc(DeviceScore, device_count);
        defer allocator.free(scores_list);

        for (scores_list) |_, i| {
            scores_list[i].has_required_exts = hasRequired(&exts_list[i]);
            scores_list[i].rt_score = rtEval(&exts_list[i]);
            scores_list[i].ext_score = extsEval(&exts_list[i]);
            scores_list[i].feat_score = featsEval(&feats_list[i]);
            scores_list[i].prop_score = propsEval(&prop_list[i]);

            const queue_support = try queues.Support.init(instance, pdevs[i], window);
            defer queue_support.deinit();

            var has_queue_support = true;
            for (queue_support.families) |family| {
                if (family == null) {
                    has_queue_support = false;
                }
            }
            scores_list[i].has_queue_support = has_queue_support;
        }

        var chosen_dev: ?usize = null;
        for (scores_list) |score, i| {
            if (!score.has_required_exts) {
                continue;
            }

            if (!score.has_queue_support) {
                continue;
            }

            if (chosen_dev) |chosen| {
                var cmp = scores_list[chosen].rt_score - score.rt_score;
                if (cmp != 0) {
                    chosen_dev = if (cmp < 0) i else chosen;
                    continue;
                }
                cmp = scores_list[chosen].ext_score - score.ext_score;
                if (cmp != 0) {
                    chosen_dev = if (cmp < 0) i else chosen;
                    continue;
                }
                cmp = scores_list[chosen].feat_score - score.feat_score;
                if (cmp != 0) {
                    chosen_dev = if (cmp < 0) i else chosen;
                    continue;
                }
                cmp = scores_list[chosen].prop_score - score.prop_score;
                if (cmp != 0) {
                    chosen_dev = if (cmp < 0) i else chosen;
                    continue;
                }
            } else {
                chosen_dev = i;
            }
        }

        if (chosen_dev) |chosen| {
            self.props = prop_list[chosen];
            self.feats = feats_list[chosen];
            self.exts = exts_list[chosen];
            self.physical_device = pdevs[chosen];
        } else unreachable;
    }

    fn hasRequired(exts: *Extensions.Device) bool {
        var has_all = false;
        has_all = has_all and exts.khr_swapchain;
        return has_all;
    }

    fn rtEval(exts: *Extensions.Device) i32 {
        var score: i32 = 0;
        if (settings.rt_on) {
            score += if (exts.khr_acceleration_structure and exts.khr_ray_tracing_pipeline) 1 else 0;
            score += if (exts.khr_ray_query) 1 else 0;
        }
        return score;
    }

    fn extsEval(exts: *Extensions.Device) i32 {
        var score: i32 = 0;
        // https://github.com/ziglang/zig/issues/137
        score += if (exts.ext_memory_budget) @as(i32, 1) else 0;
        score += if (exts.ext_hdr_metadata) @as(i32, 1) else 0;
        score += if (exts.khr_shader_float_16_int_8) @as(i32, 1) else 0;
        score += if (exts.khr_1_6bit_storage) @as(i32, 1) else 0;
        score += if (exts.khr_push_descriptor) @as(i32, 1) else 0;
        score += if (exts.ext_memory_priority) @as(i32, 1) else 0;
        score += if (exts.khr_bind_memory_2) @as(i32, 1) else 0;
        score += if (exts.khr_shader_float_controls) @as(i32, 1) else 0;
        score += if (exts.khr_spirv_1_4) @as(i32, 1) else 0;
        score += if (exts.ext_conditional_rendering) @as(i32, 1) else 0;
        score += if (exts.khr_draw_indirect_count) @as(i32, 1) else 0;
        return score;
    }

    fn featsEval(feats: *Features) i32 {
        var score: i32 = 0;

        // ------------------------------------------------------------------------
        // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkPhysicalDeviceFeatures.html

        // highly useful things
        score += if (feats.phdev.features.full_draw_index_uint_32 != 0) @as(i32, 16) else 0;
        score += if (feats.phdev.features.sampler_anisotropy != 0) @as(i32, 16) else 0;
        score += if (feats.phdev.features.texture_compression_bc != 0) @as(i32, 16) else 0;
        score += if (feats.phdev.features.independent_blend != 0) @as(i32, 16) else 0;

        // debug drawing
        score += if (feats.phdev.features.fill_mode_non_solid != 0) @as(i32, 2) else 0;
        score += if (feats.phdev.features.wide_lines != 0) @as(i32, 2) else 0;
        score += if (feats.phdev.features.large_points != 0) @as(i32, 2) else 0;

        // profiling
        // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkQueryPipelineStatisticFlagBits.html
        score += if (feats.phdev.features.pipeline_statistics_query != 0) @as(i32, 2) else 0;

        // shader features
        score += if (feats.phdev.features.fragment_stores_and_atomics != 0) @as(i32, 4) else 0;
        score += if (feats.phdev.features.shader_int_64 != 0) @as(i32, 4) else 0;
        score += if (feats.phdev.features.shader_int_16 != 0) @as(i32, 1) else 0;
        score += if (feats.phdev.features.shader_storage_image_extended_formats != 0) @as(i32, 4) else 0;

        // dynamic indexing
        score += if (feats.phdev.features.shader_uniform_buffer_array_dynamic_indexing != 0) @as(i32, 4) else 0;
        score += if (feats.phdev.features.shader_storage_buffer_array_dynamic_indexing != 0) @as(i32, 4) else 0;
        score += if (feats.phdev.features.shader_sampled_image_array_dynamic_indexing != 0) @as(i32, 4) else 0;
        score += if (feats.phdev.features.shader_storage_image_array_dynamic_indexing != 0) @as(i32, 4) else 0;
        score += if (feats.phdev.features.image_cube_array != 0) @as(i32, 1) else 0;

        // indirect and conditional rendering
        score += if (feats.phdev.features.full_draw_index_uint_32 != 0) @as(i32, 1) else 0;
        score += if (feats.phdev.features.multi_draw_indirect != 0) @as(i32, 1) else 0;
        score += if (feats.phdev.features.draw_indirect_first_instance != 0) @as(i32, 1) else 0;

        if (settings.rt_on) {
            // ------------------------------------------------------------------------
            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkPhysicalDeviceAccelerationStructureFeaturesKHR.html

            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#acceleration-structure
            score += if (feats.accstr.acceleration_structure != 0) @as(i32, 64) else 0;

            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#vkCmdBuildAccelerationStructuresIndirectKHR
            //score += if (feats.accstr.accelerationStructureIndirectBuild) 16 else 0;

            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#features-accelerationStructureHostCommands
            //score += if (feats.accstr.accelerationStructureHostCommands) 16 else 0;

            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#features-descriptorBindingAccelerationStructureUpdateAfterBind
            //score += if (feats.accstr.descriptorBindingAccelerationStructureUpdateAfterBind) 4 else 0;

            // ------------------------------------------------------------------------
            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkPhysicalDeviceRayTracingPipelineFeaturesKHR.html

            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#ray-tracing
            score += if (feats.rtpipe.ray_tracing_pipeline != 0) @as(i32, 64) else 0;

            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#vkCmdTraceRaysIndirectKHR
            //score += if (feats.rtpipe.rayTracingPipelineTraceRaysIndirect) 16 else 0;

            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#ray-traversal-culling-primitive
            //score += if (feats.rtpipe.rayTraversalPrimitiveCulling) 16 else 0;

            // ------------------------------------------------------------------------
            // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkPhysicalDeviceRayQueryFeaturesKHR.html

            // https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/KHR/SPV_KHR_ray_query.asciidoc
            score += if (feats.rquery.ray_query != 0) @as(i32, 64) else 0;
        }

        return score;
    }

    fn propsEval(props: *Props) i32 {
        var score: i32 = 0;
        score += limitsEval(&props.phdev.properties.limits);

        if (settings.rt_on) {
            score += accStrEval(props.accstr);
            score += rtPipeEval(props.rtpipe);
        }
        return score;
    }

    fn limitsEval(lims: *const vk.PhysicalDeviceLimits) i32 {
        var score: u32 = 0;
        score += log2(lims.max_image_dimension_1d);
        score += log2(lims.max_image_dimension_2d);
        score += log2(lims.max_image_dimension_3d);
        score += log2(lims.max_image_dimension_cube);
        score += log2(lims.max_image_array_layers);
        score += log2(lims.max_uniform_buffer_range);
        score += log2(lims.max_storage_buffer_range);
        score += log2(lims.max_push_constants_size);
        score += log2(lims.max_memory_allocation_count);
        score += log2(lims.max_per_stage_descriptor_storage_buffers);
        score += log2(lims.max_per_stage_descriptor_sampled_images);
        score += log2(lims.max_per_stage_descriptor_storage_images);
        score += log2(lims.max_per_stage_resources);
        score += log2(lims.max_vertex_input_attributes);
        score += log2(lims.max_vertex_input_bindings);
        score += log2(lims.max_fragment_input_components);
        score += log2(lims.max_compute_shared_memory_size);
        score += log2(lims.max_compute_work_group_invocations);
        score += log2(lims.max_draw_indirect_count);
        score += log2(lims.max_framebuffer_width);
        score += log2(lims.max_framebuffer_height);
        score += log2(lims.max_color_attachments);
        return @intCast(i32, score);
    }

    fn accStrEval(accstr: *const vk.PhysicalDeviceAccelerationStructurePropertiesKHR) u32 {
        var score: u32 = 0;
        score += log2(accstr.max_geometry_count);
        score += log2(accstr.max_instance_count);
        score += log2(accstr.max_primitive_count);
        score += log2(accstr.max_per_stage_descriptor_acceleration_structures);
        score += log2(accstr.max_descriptor_set_acceleration_structures);
        score += log2(accstr.max_descriptor_set_update_after_bind_acceleration_structures);
        return score;
    }

    fn rtPipeEval(rtpipe: *const vk.PhysicalDeviceRayTracingPipelineFeaturesKHR) u32 {
        var score: u32 = 0;
        score += log2(rtpipe.max_ray_recursion_depth);
        score += log2(rtpipe.max_ray_dispatch_invocation_count);
        score += log2(rtpipe.mayx_ray_hit_attribute_size);
        return score;
    }

    fn createDevice(self: *Self, window: *const Window, instance: *const Instance) !void {
        const queue_support = try queues.Support.init(instance, self.physical_device, window);
        defer queue_support.deinit();

        var families = std.AutoArrayHashMap(u32, i32).init(allocator);
        defer families.deinit();
        for (queue_support.families) |f| {
            if (f) |family| {
                const result = families.getOrPutAssumeCapacity(family);
                result.value_ptr.* = if (result.found_existing) result.value_ptr.* + 1 else 1;
            }
        }

        const priorities = [_]f32{ 1.0, 1.0, 1.0, 1.0 };

        var queue_infos: [queues.QueueId.count]vk.DeviceQueueCreateInfo = undefined;
        var it = families.iterator();
        {
            var i: usize = 0;
            while (it.next()) |entry| : (i += 1) {
                queue_infos[i] = .{
                    .s_type = .device_queue_create_info,
                    .flags = .{},
                    .queue_family_index = @intCast(u32, entry.key_ptr.*),
                    .queue_count = @intCast(u32, entry.value_ptr.*),
                    .p_queue_priorities = &priorities,
                };
            }
        }

        const ph_feats = self.feats.phdev.features;

        self.feats = .{
            .phdev = .{
                .s_type = .physical_device_features_2,
                .features = .{
                    .full_draw_index_uint_32 = ph_feats.full_draw_index_uint_32,
                    .sampler_anisotropy = ph_feats.sampler_anisotropy,
                    .texture_compression_bc = ph_feats.texture_compression_bc,
                    .independent_blend = ph_feats.independent_blend,

                    .fill_mode_non_solid = ph_feats.fill_mode_non_solid,
                    .wide_lines = ph_feats.wide_lines,
                    .large_points = ph_feats.large_points,

                    .fragment_stores_and_atomics = ph_feats.fragment_stores_and_atomics,
                    .shader_int_64 = ph_feats.shader_int_64,
                    .shader_int_16 = ph_feats.shader_int_16,
                    .shader_storage_image_extended_formats = ph_feats.shader_storage_image_extended_formats,

                    .shader_uniform_buffer_array_dynamic_indexing = ph_feats.shader_uniform_buffer_array_dynamic_indexing,
                    .shader_storage_buffer_array_dynamic_indexing = ph_feats.shader_storage_buffer_array_dynamic_indexing,
                    .shader_sampled_image_array_dynamic_indexing = ph_feats.shader_sampled_image_array_dynamic_indexing,
                    .shader_storage_image_array_dynamic_indexing = ph_feats.shader_storage_image_array_dynamic_indexing,
                    .image_cube_array = ph_feats.image_cube_array,
                },
            },
        };

        if (settings.rt_on) {
            self.props.accstr = .{
                .s_type = .physical_device_acceleration_structure_features_khr,
                .acceleration_structure = self.props.accstr.acceleration_structure,
                .acceleration_structure_indirect_build = self.props.accstr.acceleration_structure_indirect_build,
                .acceleration_structure_host_commands = self.props.accstr.acceleration_structure_host_commands,
            };
            self.props.rtpipe = .{
                .s_type = .physical_device_ray_tracing_pipeline_features_khr,
                .ray_tracing_pipeline = self.props.rtpipe.ray_tracing_pipeline,
                .ray_tracing_pipeline_trace_rays_indirect = self.props.rtpipe.ray_tracing_pipeline_trace_rays_indirect,
                .ray_traversal_primitive_culling = self.props.rtpipe.ray_traversal_primitive_culling,
            };
            self.props.rquery = .{
                .s_type = .physical_device_ray_query_features_khr,
                .ray_query = self.props.rquery.ray_query,
            };
        }

        if (settings.rt_on) {
            self.props.phdev.p_next = &self.props.accstr;
            self.props.accstr.p_next = &self.props.rtpipe;
            self.props.rtpipe.p_next = &self.props.rquery;
        }

        const ext_arr = Extensions.Device.toArray(allocator);
        defer allocator.free(ext_arr);

        _ = try instance.dispatch.createDevice(self.physical_device, &.{
            .flags = .{},
            .p_next = &self.props.phdev,
            .queue_create_info_count = @intCast(u32, families.count()),
            .p_queue_create_infos = &queue_infos,
            .enabled_layer_count = @intCast(u32, enabled_layers.len),
            .pp_enabled_layer_names = &enabled_layers,
            .p_enabled_features = null,
            .enabled_extension_count = @intCast(u32, ext_arr.len),
            .pp_enabled_extension_names = ext_arr.ptr,
        }, null);
    }

    fn getExtensions(avail_exts: *std.StringArrayHashMap) !std.ArrayList([][*:0]const u8) {
        const list = try std.ArrayList([*:0]const u8).init(allocator);

        for (try glfw.getRequiredDeviceExtensions()) |ext| {
            if (avail_exts.contains(ext)) {
                list.append(ext);
            }
        }

        for (Extensions.dev) |ext| {
            if (avail_exts.contains(ext)) {
                list.append(ext);
            }
        }

        return list;
    }
};