const std = @import("std");
usingnamespace @import("xrvk.zig");
const dg = @import("diligent");
const mem = std.mem;
const Allocator = mem.Allocator;

pub const Graphics = struct {
    xri: xr.InstanceDispatch,
    vki: vk.InstanceDispatch,
    vkd: vk.DeviceDispatch,

    instance: vk.Instance,
    device: vk.Device,
    physical_device: vk.PhysicalDevice,

    queue: vk.Queue,
    family_index: u32,
    command_pool: vk.CommandPool,
    swapchain: Swapchain,

    dg_device: dg.IRenderDeviceVk_Wrapper,
    dg_ctx: dg.IDeviceContext_Wrapper,

    allocator: *Allocator,

    pub fn init(
        allocator: *Allocator,
        xri: xr.InstanceDispatch,
        instance: xr.Instance,
        system: xr.SystemId,
    ) !Graphics {
        var self: Graphics = undefined;

        self.xri = xri;
        self.allocator = allocator;

        var result: xr.VkResult = undefined;

        var requirements = xr.GraphicsRequirementsVulkan2KHR.empty();
        try xri.getVulkanGraphicsRequirements2KHR(instance, system, &requirements);

        // CREATE INSTANCE

        const app_info = vk.ApplicationInfo{
            .p_application_name = "test",
            .application_version = vk.makeVersion(0, 0, 0),
            .p_engine_name = "test",
            .engine_version = vk.makeVersion(0, 0, 0),
            .api_version = vk.API_VERSION_1_2,
        };

        const zero = [_:0]u8{};
        const instance_extensions = [_][*:0]const u8{
            "VK_EXT_debug_utils",
            "VK_KHR_get_physical_device_properties2",
            "VK_KHR_surface",
        };

        const instance_create_info = vk.InstanceCreateInfo{
            .flags = .{},
            .p_application_info = &app_info,
            .enabled_layer_count = 0,
            .pp_enabled_layer_names = @ptrCast([*]const [*:0]const u8, &zero),
            .enabled_extension_count = instance_extensions.len,
            .pp_enabled_extension_names = &instance_extensions,
        };

        try xri.createVulkanInstanceKHR(instance, .{
            .create_flags = .{},
            .system_id = system,
            .pfn_get_instance_proc_addr = vk.getInstanceProcAddr,
            .vulkan_create_info = &instance_create_info,
            .vulkan_allocator = null,
        }, &self.instance, &result);
        switch (result) {
            .success => {},
            else => {
                std.debug.print("failed to create VkInstance: {}\n", .{result});
                return error.InstanceCreateFailed;
            },
        }
        self.vki = try vk.InstanceDispatch.load(self.instance, vk.getInstanceProcAddr);
        errdefer self.vki.destroyInstance(self.instance, null);

        // CREATE PHYSICAL DEVICE

        try xri.getVulkanGraphicsDevice2KHR(instance, .{
            .system_id = system,
            .vulkan_instance = self.instance,
        }, &self.physical_device);

        // CHOOSE QUEUE

        var family_count: u32 = undefined;
        self.vki.getPhysicalDeviceQueueFamilyProperties(self.physical_device, &family_count, null);

        const families = try allocator.alloc(vk.QueueFamilyProperties, family_count);
        defer allocator.free(families);
        self.vki.getPhysicalDeviceQueueFamilyProperties(
            self.physical_device,
            &family_count,
            families.ptr,
        );

        self.family_index = for (families) |properties, i| {
            const family = @intCast(u32, i);

            if (properties.queue_flags.contains(.{ .graphics_bit = true })) {
                break family;
            }
        } else {
            return error.NoGraphicsFamily;
        };

        const priority = [_]f32{0};
        const queue_create_info = [_]vk.DeviceQueueCreateInfo{
            .{
                .flags = .{},
                .queue_count = 1,
                .p_queue_priorities = &priority,
                .queue_family_index = self.family_index,
            },
        };

        // CREATE DEVICE

        const device_extensions = [_][*:0]const u8{
            // "VK_KHR_maintenance1",
            "VK_KHR_8bit_storage",
            "VK_KHR_16bit_storage",
            "VK_KHR_storage_buffer_storage_class",
        };

        const device_features = mem.zeroInit(vk.PhysicalDeviceFeatures, .{
            .geometry_shader = vk.TRUE,
            .multi_viewport = vk.TRUE,
        });

        const device_create_info = vk.DeviceCreateInfo{
            .flags = .{},
            .queue_create_info_count = queue_create_info.len,
            .p_queue_create_infos = &queue_create_info,
            .enabled_layer_count = 0,
            .pp_enabled_layer_names = undefined,
            .enabled_extension_count = device_extensions.len,
            .pp_enabled_extension_names = device_extensions[0..],
            .p_enabled_features = &device_features,
        };

        try xri.createVulkanDeviceKHR(instance, .{
            .system_id = system,
            .create_flags = .{},
            .pfn_get_instance_proc_addr = vk.getInstanceProcAddr,
            .vulkan_physical_device = self.physical_device,
            .vulkan_create_info = &device_create_info,
            .vulkan_allocator = null,
        }, &self.device, &result);
        switch (result) {
            .success => {},
            else => {
                std.debug.print("failed to create VkDevice: {}\n", .{result});
                return error.DeviceCreateFailed;
            },
        }
        self.vkd = try vk.DeviceDispatch.load(self.device, vk.getDeviceProcAddr);
        errdefer self.vkd.destroyDevice(self.device, null);

        // GET QUEUE

        self.queue = self.vkd.getDeviceQueue(self.device, self.family_index, 0);

        // DILIGENT

        const create_info = dg.EngineVkCreateInfo{
            ._EngineCreateInfo = .{
                .APIVersion = dg.DILIGENT_API_VERSION,
                .NumDeferredContexts = 0,
                .Features = mem.zeroInit(dg.DeviceFeatures, .{
                    .GeometryShaders = dg.DEVICE_FEATURE_STATE_ENABLED,
                    .MultiViewport = dg.DEVICE_FEATURE_STATE_ENABLED,
                }),
                .pRawMemAllocator = null,
                .DebugMessageCallback = null,
            },
            .NumCommandsToFlushCmdBuffer = 2048,
            .MainDescriptorPoolSize = .{
                .MaxDescriptorSets = 8192,
                .NumSeparateSamplerDescriptors = 1024,
                .NumCombinedSamplerDescriptors = 8192,
                .NumSampledImageDescriptors = 8192,
                .NumStorageImageDescriptors = 1024,
                .NumUniformBufferDescriptors = 4096,
                .NumStorageBufferDescriptors = 4096,
                .NumUniformTexelBufferDescriptors = 1024,
                .NumStorageTexelBufferDescriptors = 1024,
                .NumInputAttachmentDescriptors = 256,
                .NumAccelStructDescriptors = 256,
            },
            .DynamicDescriptorPoolSize = .{
                .MaxDescriptorSets = 2048,
                .NumSeparateSamplerDescriptors = 256,
                .NumCombinedSamplerDescriptors = 2048,
                .NumSampledImageDescriptors = 2048,
                .NumStorageImageDescriptors = 256,
                .NumUniformBufferDescriptors = 1024,
                .NumStorageBufferDescriptors = 1024,
                .NumUniformTexelBufferDescriptors = 256,
                .NumStorageTexelBufferDescriptors = 256,
                .NumInputAttachmentDescriptors = 64,
                .NumAccelStructDescriptors = 64,
            },
            .DeviceLocalMemoryPageSize = 16 << 20,
            .HostVisibleMemoryPageSize = 16 << 20,
            .DeviceLocalMemoryReserveSize = 256 << 20,
            .HostVisibleMemoryReserveSize = 256 << 20,
            .UploadHeapPageSize = 1 << 20,
            .DynamicHeapSize = 8 << 20,
            .DynamicHeapPageSize = 256 << 10,
            .QueryPoolSizes = .{ 0, 128, 128, 512, 128, 256 },
            .pDxCompilerPath = null,
            .AdapterId = 0,
            .EnableValidation = false,
            // .GlobalExtensionCount = @intCast(u32, instance_extensions.len),
            // .ppGlobalExtensionNames = instance_extensions,
            .GlobalExtensionCount = 0,
            .ppGlobalExtensionNames = null,
            .pVkAllocator = null,
        };

        const factory = dg.IEngineFactoryVk_Wrapper.wrap(dg.Diligent_GetEngineFactoryVk());

        var device: ?*dg.IRenderDevice = null;
        var context: ?*dg.IDeviceContext = null;
        factory.AttachToVulkanDevice(
            &@intToPtr(*dg.VkInstance_T, @enumToInt(self.instance)),
            &@intToPtr(*dg.VkPhysicalDevice_T, @enumToInt(self.physical_device)),
            @ptrCast(*const dg.VkPhysicalDeviceFeatures, &device_features),
            &@intToPtr(*dg.VkDevice_T, @enumToInt(self.device)),
            &@intToPtr(*dg.VkQueue_T, @enumToInt(self.queue)),
            self.family_index,
            &create_info,
            &device,
            &context,
        );

        self.dg_device = dg.IRenderDeviceVk_Wrapper.wrap(@ptrCast(?*dg.IRenderDeviceVk, device));
        self.dg_ctx = dg.IDeviceContext_Wrapper.wrap(context);

        return self;
    }

    pub fn createSwapchain(self: *Graphics, session: xr.Session, configs: []const xr.ViewConfigurationView) !void {
        const factory = dg.IEngineFactoryVk_Wrapper.wrap(dg.Diligent_GetEngineFactoryVk());

        try Swapchain.init(&self.swapchain, self, factory, session, configs);
    }

    pub fn deinit(self: *const Graphics) void {
        self.swapchain.deinit();

        self.vkd.destroyCommandPool(self.device, self.command_pool, null);
        self.vkd.destroyDevice(self.device, null);
        self.vki.destroyInstance(self.instance, null);

        self.allocator.destroy(self);
    }

    pub fn lockQueue(self: *const Graphics) void {
        _ = self.dg_device.LockCommandQueue(0);
    }

    pub fn unlockQueue(self: *const Graphics) void {
        self.dg_device.UnlockCommandQueue(0);
    }

    pub fn setupDebugging(self: *Graphics) !void {
        if (self.vki.vkCreateDebugUtilsMessengerEXT) |pfn| {
            var messenger: vk.c.VkDebugUtilsMessengerEXT = undefined;
            const messenger_create_info = vk.c.VkDebugUtilsMessengerCreateInfoEXT{
                .sType = .VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT,
                .pNext = null,
                .flags = 0,
                .messageSeverity = 0x1110,
                .messageType = 0x3,
                .pfnUserCallback = @This().debug_cb,
                .pUserData = @as(*c_void, self),
            };

            _ = pfn(
                @intToPtr(vk.c.VkInstance, @enumToInt(self.instance)),
                &messenger_create_info,
                null,
                &messenger,
            );
        } else {
            return error.FunctionUnsupported;
        }
    }

    pub fn getBinding(self: *const Graphics) xr.GraphicsBindingVulkan2KHR {
        return .{
            .instance = self.instance,
            .device = self.device,
            .physical_device = self.physical_device,
            .queue_family_index = self.family_index,
            .queue_index = 0,
        };
    }

    fn debug_cb(
        msg_severity: vk.c.VkDebugUtilsMessageSeverityFlagBitsEXT,
        msg_type: vk.c.VkDebugUtilsMessageTypeFlagsEXT,
        c_data: ?*const vk.c.VkDebugUtilsMessengerCallbackDataEXT,
        c_self: ?*c_void,
    ) callconv(.C) vk.Bool32 {
        const level = if (msg_severity == .VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT)
            "error"
        else
            "warn";
        const self = @ptrCast(*Graphics, @alignCast(@alignOf(Graphics), c_self));

        if (c_data) |data| {
            std.debug.print("[{}] {}\n", .{
                level,
                @as([*:0]const u8, data.pMessage),
            });
        }
        return 0;
    }
};

pub const Swapchain = struct {
    handle: xr.Swapchain,
    dg_handle: dg.ISwapChainVk_Wrapper,
    image_rects: []xr.Rect2Di,

    graphics: *const Graphics,

    pub fn init(
        self: *Swapchain,
        gfx: *const Graphics,
        factory: dg.IEngineFactoryVk_Wrapper,
        session: xr.Session,
        configs: []const xr.ViewConfigurationView,
    ) !void {
        const config = configs[0];
        const image_rects = try gfx.allocator.alloc(xr.Rect2Di, configs.len);
        errdefer gfx.allocator.free(image_rects);

        for (configs) |c, i| {
            if (c.recommended_swapchain_sample_count != config.recommended_swapchain_sample_count) {
                return error.ViewCfgDiffer;
            }
            if (c.recommended_image_rect_width != config.recommended_image_rect_width) {
                return error.ViewCfgDiffer;
            }
            if (c.recommended_image_rect_height != config.recommended_image_rect_height) {
                return error.ViewCfgDiffer;
            }

            image_rects[i] = .{
                .offset = .{
                    .y = @intCast(i32, config.recommended_image_rect_height) * @intCast(i32, i),
                },
                .extent = .{
                    .width = @intCast(i32, config.recommended_image_rect_width),
                    .height = @intCast(i32, config.recommended_image_rect_height),
                },
            };
        }

        // xri.enumerateSwapchainFormats(...);
        const view_count = @intCast(u32, configs.len);
        const swapchain = try gfx.xri.createSwapchain(session, .{
            .create_flags = .{},
            .usage_flags = .{
                .color_attachment_bit = true,
                .transfer_dst_bit = true,
                .sampled_bit = true,
            },
            .format = @enumToInt(vk.Format.b8g8r8a8_srgb),
            .sample_count = config.recommended_swapchain_sample_count,
            .width = config.recommended_image_rect_width,
            .height = config.recommended_image_rect_height * view_count,
            .face_count = 1,
            .array_size = 1,
            .mip_count = 1,
        });

        // get images
        const image_count = try gfx.xri.enumerateSwapchainImages(swapchain, 0, null);
        const xr_images = try gfx.allocator.alloc(xr.SwapchainImageVulkan2KHR, image_count);
        mem.set(xr.SwapchainImageVulkan2KHR, xr_images, xr.SwapchainImageVulkan2KHR.empty());
        defer gfx.allocator.free(xr_images);

        _ = try gfx.xri.enumerateSwapchainImages(swapchain, image_count, @ptrCast(
            [*]xr.SwapchainImageBaseHeader,
            xr_images.ptr,
        ));

        const images = try gfx.allocator.alloc(vk.Image, xr_images.len);
        defer gfx.allocator.free(images);

        for (xr_images) |xr_image, i| {
            images[i] = xr_image.image;
        }

        const swap_desc = dg.SwapChainDesc{
            .Width = config.recommended_image_rect_width,
            .Height = config.recommended_image_rect_height * view_count,
            .ColorBufferFormat = dg.TEX_FORMAT_BGRA8_UNORM_SRGB,
            .DepthBufferFormat = dg.TEX_FORMAT_D32_FLOAT,
            .Usage = dg.SWAP_CHAIN_USAGE_RENDER_TARGET,
            .PreTransform = dg.SURFACE_TRANSFORM_OPTIMAL,
            .BufferCount = @intCast(u32, images.len),
            .DefaultDepthValue = 1,
            .DefaultStencilValue = 0,
            .IsPrimary = true,
        };
        const callbacks = dg.SwapChainImageCallbacks{
            .Cookie = self,
            .ImageCallbackAcquire = @This().acquire_cb,
            .ImageCallbackRelease = @This().release_cb,
        };

        var dg_swapchain: ?*dg.ISwapChain = undefined;
        factory.CreateSwapChainVk(
            @ptrCast(*dg.IRenderDevice, gfx.dg_device.ptr),
            gfx.dg_ctx.ptr,
            &swap_desc,
            @intCast(u32, images.len),
            @ptrCast(**dg.VkImage_T, images.ptr),
            &callbacks,
            &dg_swapchain,
        );

        std.debug.print("created swapchain {}x{}*{}, SS{}, {} images\n", .{
            config.recommended_image_rect_width,
            config.recommended_image_rect_height,
            view_count,
            config.recommended_swapchain_sample_count,
            images.len,
        });

        self.* = .{
            .handle = swapchain,
            .dg_handle = dg.ISwapChainVk_Wrapper.wrap(@ptrCast([*c]dg.ISwapChainVk, dg_swapchain)),
            .image_rects = image_rects,

            .graphics = gfx,
        };
    }

    pub fn deinit(self: *const Swapchain) void {
        const gfx = self.graphics;

        gfx.allocator.free(self.image_rects);
        gfx.xri.destroySwapchain(self.handle) catch {};
    }

    pub fn acquireImage(self: *const Swapchain) !void {
        const result = self.dg_handle.AcquireNextImage();
        if (result != .VK_SUCCESS) {
            return error.Unknown;
        }
    }

    pub fn presentImage(self: *const Swapchain) void {
        self.dg_handle.PresentImage();
    }

    fn acquire_cb(c_self: ?*c_void, index: [*c]u32) callconv(.C) dg.VkResult {
        const self = @ptrCast(*Swapchain, @alignCast(@alignOf(Swapchain), c_self));
        const xri = &self.graphics.xri;

        index.* = xri.acquireSwapchainImage(
            self.handle,
            &xr.SwapchainImageAcquireInfo.empty(),
        ) catch return .VK_ERROR_UNKNOWN;

        _ = xri.waitSwapchainImage(self.handle, .{ .timeout = -1 }) catch {};

        return .VK_SUCCESS;
    }

    fn release_cb(c_self: ?*c_void) callconv(.C) void {
        const self = @ptrCast(*Swapchain, @alignCast(@alignOf(Swapchain), c_self));

        _ = self.graphics.xri.releaseSwapchainImage(
            self.handle,
            &xr.SwapchainImageReleaseInfo.empty(),
        ) catch {};
    }
};