#include <cassert>
#include <cstdio>
#include <ratio>
#include <span>
#include <array>
#include <chrono>
#include <vk_mem_alloc.h>

#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>

#include "imgui.h"

#include "handle.hpp"
#include "window.hpp"
#include "vk-device.hpp"
#include "utils.hpp"

struct device_texture;

#ifdef _DEBUG
#define ENABLE_RENDERDOC
#endif // _DEBUG

#ifdef ENABLE_RENDERDOC
#include <renderdoc.h>

#define RD_START_CAPTURE if(rdoc_api != nullptr) rdoc_api->StartFrameCapture(nullptr, nullptr)
#define RD_END_CAPTURE if(rdoc_api != nullptr) rdoc_api->EndFrameCapture(nullptr, nullptr)

RENDERDOC_API_1_4_1 *rdoc_api = nullptr;
#else
#define RD_START_CAPTURE (0)
#define RD_END_CAPTURE (0)
#endif // ENABLE_RENDERDOC

static constexpr size_t Kb = 1024U;
static constexpr size_t Mb = 1024U * Kb;

void ui(uint32_t frame_time) {
    static float max_fps = 0.f;
    static uint32_t max_frame_time = 1000U;
    float fps = 1000000U / frame_time;
    if (fps > max_fps) {
        max_fps = fps;
    }
    if (frame_time > max_frame_time) {
        max_frame_time = frame_time;
    }
    ImGui::NewFrame();
    // bool open = true;
    // ImGui::ShowDemoWindow(&open);
    ImGui::Text("FPS : %f\n", fps);
    ImGui::Text("Max FPS : %f\n", max_fps);
    ImGui::Text("frame_time : %f ms\n", frame_time / 1000.f);
    ImGui::Text("max frame_time : %f ms\n", max_frame_time / 1000.f);
    ImGui::EndFrame();
    ImGui::Render();
}

void update_buffers(ImDrawData* draw_data, const device_buffer& vertex_buffer, const device_buffer& index_buffer) {
    off_t vertex_offset = 0;
    off_t index_offset = 0;
    for (auto index {0} ; index < draw_data->CmdListsCount; index++) {
        auto* cmd_list = draw_data->CmdLists[index];

        size_t vertex_bytes_size = sizeof(ImDrawVert) * cmd_list->VtxBuffer.size();
        auto* vtx_ptr = static_cast<uint8_t*>(vertex_buffer.mapped_ptr) + vertex_offset;
        memcpy(vtx_ptr, cmd_list->VtxBuffer.Data, vertex_bytes_size);
        vertex_offset += static_cast<off_t>(vertex_bytes_size);

        size_t index_bytes_size = sizeof(ImDrawIdx) * cmd_list->IdxBuffer.size();
        memcpy(static_cast<uint8_t*>(index_buffer.mapped_ptr) + index_offset, cmd_list->IdxBuffer.Data, index_bytes_size);
        index_offset += static_cast<off_t>(index_bytes_size);
    }
}

int main() {

#ifdef ENABLE_RENDERDOC
    if (HMODULE mod = GetModuleHandleA("renderdoc.dll")) {
        auto RENDERDOC_GetAPI = (pRENDERDOC_GetAPI)GetProcAddress(mod, "RENDERDOC_GetAPI");
        assert(RENDERDOC_GetAPI(eRENDERDOC_API_Version_1_1_2, (void **)&rdoc_api) == 1);
    }
#endif

    constexpr size_t window_width = 1280U;
    constexpr size_t window_height = 720U;
    window wnd { window_width, window_height };

    ImGui::CreateContext();
    ImGuiIO &io = ImGui::GetIO();
    io.DisplaySize.x = window_width;
    io.DisplaySize.y = window_height;

    auto& device = vkdevice::get_render_device();
    device.init();

    auto surface_handle = device.create_surface({
        .window_handle = wnd.handle,
        .surface_format = {
            .format = VK_FORMAT_B8G8R8A8_UNORM,
            .colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
        },
        .present_mode = VK_PRESENT_MODE_IMMEDIATE_KHR,
        .usages = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
    });
    auto& surface = device.get_surface(surface_handle);

    constexpr size_t image_size = 1024U;
    constexpr size_t image_count = 4U;
    auto staging_buffer_handle = device.create_buffer({
        .size = image_size * image_size * sizeof(uint32_t) * image_count,
        .usages = VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
        .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
        .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
    });
    const auto& staging_buffer = device.get_buffer(staging_buffer_handle);

    int width;
    int height;
    int channels;
    uint8_t* data = stbi_load("../models/sponza/466164707995436622.jpg", &width, &height, &channels, 4);
    size_t size = static_cast<size_t>(width) * height * 4U;
    memcpy(staging_buffer.mapped_ptr, data, size);

    auto gpu_texture_handle = device.create_texture({
        .width  = static_cast<uint32_t>(width),
        .height = static_cast<uint32_t>(height),
        .usages = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
        .format = VK_FORMAT_R8G8B8A8_UNORM,
        .type   = VK_IMAGE_TYPE_2D,
    });

    off_t offset = width * height * 4;
    data = stbi_load("../models/sponza/715093869573992647.jpg", &width, &height, &channels, 4);
    size = static_cast<size_t>(width) * height * 4U;
    auto* addr = static_cast<void*>(static_cast<uint8_t*>(staging_buffer.mapped_ptr) + offset);
    memcpy(addr, data, size);

    auto second_texture = device.create_texture({
        .width  = static_cast<uint32_t>(width),
        .height = static_cast<uint32_t>(height),
        .usages = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_STORAGE_BIT,
        .format = VK_FORMAT_R8G8B8A8_UNORM,
        .type   = VK_IMAGE_TYPE_2D,
    });

    auto output_texture_handle = device.create_texture({
        .width  = static_cast<uint32_t>(width),
        .height = static_cast<uint32_t>(height),
        .usages = VK_IMAGE_USAGE_STORAGE_BIT,
        .format = VK_FORMAT_R8G8B8A8_UNORM,
        .type   = VK_IMAGE_TYPE_2D,
    });

    uint8_t *pixels = nullptr;
    int atlas_width;
    int atlas_height;
    io.Fonts->GetTexDataAsRGBA32(&pixels, &atlas_width, &atlas_height);
    const auto ui_texture_handle = device.create_texture({
        .width = static_cast<uint32_t>(atlas_width),
        .height = static_cast<uint32_t>(atlas_height),
        .usages = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
        .format = VK_FORMAT_R8G8B8A8_UNORM,
        .type = VK_IMAGE_TYPE_2D,
    });
    const auto& ui_texture = device.get_texture(ui_texture_handle);
    io.Fonts->SetTexID(*(void **)&ui_texture_handle.id);

    off_t ui_offset = 2L * offset;
    size = static_cast<size_t>(atlas_width) * atlas_height * 4U;
    addr = static_cast<void*>(static_cast<uint8_t*>(staging_buffer.mapped_ptr) + ui_offset);
    memcpy(addr, pixels, size);

    auto vertex_buffer_handle = device.create_buffer({
        .size = Mb * sizeof(uint32_t),
        .usages = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
        .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
        .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
    });
    auto vertex_buffer = device.get_buffer(vertex_buffer_handle);

    auto index_buffer_handle = device.create_buffer({
        .size = Mb * sizeof(uint32_t),
        .usages = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT,
        .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
        .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
    });
    auto index_buffer = device.get_buffer(index_buffer_handle);

    auto code = read_file("shaders/test.comp.spv");
    auto compute = device.create_pipeline({
        .cs_code = code,
    });

    std::array<VkFormat, 1U> formats { surface.surface_format.format };
    auto vertex_code = read_file("shaders/api-test-ui.vert.spv");
    auto fragment_code = read_file("shaders/api-test-ui.frag.spv");
    auto graphics = device.create_pipeline({
        .vs_code = vertex_code,
        .fs_code = fragment_code,
        .color_attachments_format = formats
    });

    auto semaphore = device.create_semaphore({
        .type = VK_SEMAPHORE_TYPE_TIMELINE,
    });
    constexpr uint32_t virtual_frames_count = 2U;
    handle<device_semaphore> submit_semaphores[virtual_frames_count];
    for (auto& submit_semaphore : submit_semaphores) {
        submit_semaphore = device.create_semaphore({
            .type = VK_SEMAPHORE_TYPE_BINARY,
        });
    }
    handle<device_semaphore> acquire_semaphores[virtual_frames_count];
    for (auto& acquire_semaphore : acquire_semaphores) {
        acquire_semaphore = device.create_semaphore({
            .type = VK_SEMAPHORE_TYPE_BINARY,
        });
    }

    std::array<graphics_command_buffer, 4U> graphics_command_buffers;
    device.allocate_command_buffers(graphics_command_buffers, QueueType::GRAPHICS);

    RD_START_CAPTURE;

    graphics_command_buffers[0].start();

    graphics_command_buffers[0].barrier(gpu_texture_handle, VK_PIPELINE_STAGE_2_TRANSFER_BIT, VK_ACCESS_2_TRANSFER_WRITE_BIT, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    graphics_command_buffers[0].barrier(second_texture, VK_PIPELINE_STAGE_2_TRANSFER_BIT, VK_ACCESS_2_TRANSFER_WRITE_BIT, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    graphics_command_buffers[0].barrier(ui_texture_handle, VK_PIPELINE_STAGE_2_TRANSFER_BIT, VK_ACCESS_2_TRANSFER_WRITE_BIT, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    graphics_command_buffers[0].copy(staging_buffer_handle, gpu_texture_handle);
    graphics_command_buffers[0].copy(staging_buffer_handle, second_texture, static_cast<VkDeviceSize>(offset));
    graphics_command_buffers[0].copy(staging_buffer_handle, ui_texture_handle, static_cast<VkDeviceSize>(ui_offset));

    graphics_command_buffers[0].barrier(gpu_texture_handle, VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_WRITE_BIT, VK_IMAGE_LAYOUT_READ_ONLY_OPTIMAL);
    graphics_command_buffers[0].barrier(second_texture, VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, VK_ACCESS_2_SHADER_READ_BIT, VK_IMAGE_LAYOUT_READ_ONLY_OPTIMAL);

    struct compute_params {
        uint32_t input1;
        uint32_t input2;
        uint32_t output;
    };
    const auto& uniform_buffer = device.get_bindingmodel().get_uniform_buffer();
    compute_params params {
        device.get_texture(second_texture).get_storage_index(),
        device.get_texture(gpu_texture_handle).get_storage_index(),
        device.get_texture(output_texture_handle).get_storage_index(),
    };
    memcpy(uniform_buffer.mapped_ptr, &params, sizeof(params));

    graphics_command_buffers[0].dispatch({
        .pipeline = compute,
        .group_size = { .vec = { image_size, image_size, 1U }},
        .local_group_size = { .vec = { 8U, 8U, 1U }},
        .uniforms_offset = 0U,
    });

    graphics_command_buffers[0].barrier(gpu_texture_handle, VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT, VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);

    graphics_command_buffers[0].stop();

    device.submit(std::span { static_cast<command_buffer*>(graphics_command_buffers.data()), 1 }, {}, semaphore);

    RD_END_CAPTURE;

    // std::array<struct command_buffer, 1U> ui_command_buffers = {command_buffers[1]};
    std::array<handle<device_texture>, 1U> color_attachment {};
    uint32_t frame_time;
    uint32_t frame_count = 0U;
    while(wnd.isOpen) {
        auto start = std::chrono::system_clock::now();
        wnd.poll_events();

        uint32_t virtual_frame_index = frame_count % virtual_frames_count;
        auto& submit_semaphore = submit_semaphores[virtual_frame_index];
        auto& acquire_semaphore = acquire_semaphores[virtual_frame_index];
        auto command_buffers = std::span<struct command_buffer> { static_cast<struct command_buffer*>(graphics_command_buffers.data() + virtual_frame_index + 1U), 1 };
        graphics_command_buffer& command_buffer = graphics_command_buffers[virtual_frame_index + 1U];
        for (auto& event : wnd.events) {
            if (event.type == EVENT_TYPES::MOUSE_MOVE) {
                io.MousePos.x = static_cast<float>(event.x);
                io.MousePos.y = static_cast<float>(event.y);
            }

            io.MouseDown[0] = event.type == EVENT_TYPES::BUTTON_PRESS;
        }

        ui(frame_time);

        auto* draw_data = ImGui::GetDrawData();
        update_buffers(draw_data, vertex_buffer, index_buffer);

        // device.wait(semaphore);
        device.wait(submit_semaphore);

        surface.acquire_image_index(acquire_semaphore);
        color_attachment[0] = surface.get_backbuffer_image();

        RD_START_CAPTURE;
        command_buffer.start(),

        command_buffer.barrier(surface.get_backbuffer_image(), VK_PIPELINE_STAGE_2_TOP_OF_PIPE_BIT, VK_ACCESS_2_NONE, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);
        command_buffer.begin_renderpass({
            .render_area = {
                .x = static_cast<int32_t>(draw_data->DisplayPos.x),
                .y = static_cast<int32_t>(draw_data->DisplayPos.y),
                .width = static_cast<uint32_t>(draw_data->DisplaySize.x),
                .height = static_cast<uint32_t>(draw_data->DisplaySize.y),
            },
            .color_attachments = color_attachment,
        });

        struct draw_params {
            uintptr_t vertex_buffer;
            float scale[2U];
            float translate[2U];
            uint32_t texture_index;
        };

        uint32_t draw_index = 0U;
        std::array<draw_params, 100U> draws;
        for (auto index {0} ; index < draw_data->CmdListsCount; index++) {
            auto* cmd_list = draw_data->CmdLists[index];

            for (auto& draw_command : cmd_list->CmdBuffer) {
                draws[draw_index] = draw_params{
                    .vertex_buffer = vertex_buffer.device_address,
                    .scale {
                        2.f / draw_data->DisplaySize.x,
                        2.f / draw_data->DisplaySize.y,
                    },
                    .translate {
                        -1.f - draw_data->DisplayPos.x * (2.f / draw_data->DisplaySize.x),
                        -1.f - draw_data->DisplayPos.y * (2.f / draw_data->DisplaySize.y),
                    },
                    .texture_index = ui_texture.get_sampled_index(),
                };

                command_buffer.render({
                    .pipeline = graphics,
                    .index_buffer = index_buffer_handle,
                    .vertex_count = draw_command.ElemCount,
                    .vertex_offset = draw_command.VtxOffset,
                    .index_offset = draw_command.IdxOffset,
                    .instance_count = 1U,
                    .uniforms_offset = 0U,
                });

                draw_index++;
            }
        }

        memcpy(uniform_buffer.mapped_ptr, draws.data(), sizeof(draw_params) * draw_index);

        command_buffer.end_renderpass();
        command_buffer.barrier(surface.get_backbuffer_image(), VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT, VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT, VK_IMAGE_LAYOUT_PRESENT_SRC_KHR);

        command_buffer.stop();

        // device.submit(command_buffers, acquire_semaphore, semaphore);
        // device.wait(semaphore);
        // device.present(surface_handle, {});

        device.submit(command_buffers, acquire_semaphore, submit_semaphore);
        device.present(surface_handle, submit_semaphore);

        auto end = std::chrono::system_clock::now();
        frame_time = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();

        RD_END_CAPTURE;
    }

    device.wait(semaphore);

    device.destroy_surface(surface_handle);

    device.destroy_pipeline(compute);
    device.destroy_pipeline(graphics);

    device.destroy_semaphore(semaphore);
    for (auto& submit_semaphore : submit_semaphores) {
        device.destroy_semaphore(submit_semaphore);
    }
    for (auto& acquire_semaphore : acquire_semaphores) {
        device.destroy_semaphore(acquire_semaphore);
    }

    device.destroy_texture(gpu_texture_handle);
    device.destroy_texture(second_texture);
    device.destroy_texture(output_texture_handle);
    device.destroy_texture(ui_texture_handle);

    device.destroy_buffer(staging_buffer_handle);
    device.destroy_buffer(vertex_buffer_handle);
    device.destroy_buffer(index_buffer_handle);

    ImGui::DestroyContext();

    return 0;
}