//
// Created by Vicente Ferrari Smith on 13.02.26.
//

#ifndef V_RENDERER_H
#define V_RENDERER_H

#include "init.h"
#include <volk/volk.h>
#include <GLFW/glfw3.h>
#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#define GLM_ENABLE_EXPERIMENTAL
#include <glm/glm.hpp>
#include <glm/ext/matrix_clip_space.hpp>
#include "glm/gtx/string_cast.hpp"
#include <vma/vk_mem_alloc.h>
#include "sprite.h"
#include "texture.h"
#include <misc.h>
#include <array>
#include <span>
#include <slang/slang.h>
#include <slang/slang-com-ptr.h>

inline Slang::ComPtr<slang::IGlobalSession> slangGlobalSession;

enum class PROJECTION_TYPE : uint8_t {
    NONE,
    ORTHOGRAPHIC_WORLD,
    ORTHOGRAPHIC_WINDOW,
    PERSPECTIVE_WORLD,
    PERSPECTIVE_WINDOW,
    COUNT,
};

struct vertex_p2_s2_st2_col4_a1_u32 {
    glm::vec2 pos;
    glm::vec2 scale;
    glm::vec2 uv;
    glm::vec4 color;
    float     alpha;
    uint32_t  textureID;

    static VkVertexInputBindingDescription getBindingDescription() {
        return {0, sizeof(vertex_p2_s2_st2_col4_a1_u32), VK_VERTEX_INPUT_RATE_VERTEX};
    }

    static std::array<VkVertexInputAttributeDescription, 6> getAttributeDescriptions() {
        return {
            {
                {0, 0, VK_FORMAT_R32G32_SFLOAT, offsetof(vertex_p2_s2_st2_col4_a1_u32, pos)},
                {1, 0, VK_FORMAT_R32G32_SFLOAT, offsetof(vertex_p2_s2_st2_col4_a1_u32, scale)},
                {2, 0, VK_FORMAT_R32G32_SFLOAT, offsetof(vertex_p2_s2_st2_col4_a1_u32, uv)},
                {3, 0, VK_FORMAT_R32G32B32A32_SFLOAT, offsetof(vertex_p2_s2_st2_col4_a1_u32, color)},
                {4, 0, VK_FORMAT_R32_SFLOAT, offsetof(vertex_p2_s2_st2_col4_a1_u32, alpha)},
                {5, 0, VK_FORMAT_R32_UINT, offsetof(vertex_p2_s2_st2_col4_a1_u32, textureID)},
            }
        };
    }
};

// commands

enum class PipelineType : uint8_t {
    None,
    TexturedQuad,
    ColoredQuad,
    Line,
    Text,
    Chunk
};

struct TexturedQuadCmd {
    glm::vec2 position;
    glm::vec2 size;
    glm::vec2 uvMin;
    glm::vec2 uvMax;
    glm::vec4 color;
    uint16_t textureID;
};

struct ColoredQuadCmd {
    glm::vec2 pos;
    glm::vec2 scale;
    glm::vec4 color;
};

struct LineCmd {
    glm::vec2 start;
    glm::vec2 end;
    glm::vec4 color;
};

// struct TextCmd {
//     Font* font;
//     std::string text;
//     glm::vec2 position;
//     glm::vec4 color;
// };

struct ChunkCmd {
    VkBuffer vertexBuffer;
    VkBuffer indexBuffer;
    uint32_t indexCount;
};

struct SortKey {
    uint16_t depth; // world Z or Y-sorted depth
    uint16_t materialID; // texture sheet, font atlas, etc.
    uint8_t pipeline; // PipelineType

    bool operator<(const SortKey& b) const;
};

struct RenderCommand {

    SortKey      key;
    PipelineType pipeline;

    union {
        TexturedQuadCmd textured_quad;
        ColoredQuadCmd  colored_quad;
        LineCmd         line;
        // TextCmd         text;
        ChunkCmd        chunk;
    };
};

////////////////////////////////////////////////////////////////////////////////////////////////

struct AllocatedBuffer {
    VkBuffer buffer;
    VmaAllocation allocation;
    VmaAllocationInfo info;
};

struct GPUMeshBuffers {
    AllocatedBuffer indexBuffer;
    AllocatedBuffer vertexBuffer;
    VkDeviceAddress vertexBufferAddress;
};

struct Renderer {
    std::vector<RenderCommand> commands{};

    VkDescriptorSetLayout descriptor_set_layout{};
    VkPipelineLayout pipelineLayout{};
    VkPipeline textured_quad_pipeline{};
    VkPipeline colored_quad_pipeline{};
    VkPipeline line_pipeline{};
    VkPipeline text_pipeline{};
    VkPipeline chunk_pipeline{};
    VkDescriptorSet set{};

    VkSampler defaultSampler{};

    uint32_t nextTextureSlot = 0;

    struct Frame {
        VkCommandPool commandPool{};
        VkCommandBuffer command_buffer{};

        VkSemaphore imageAvailable{};
        VkFence in_flight_fence{};

        AllocatedBuffer vertexBuffer{};
    };

    std::vector<Frame> frames;
    uint32_t currentFrame = 0;

    VkDescriptorPool descriptorPool{};
    std::vector<VkDescriptorSet> textureSets{};

    void begin_frame();
    void end_frame();
    void flush();

    void submit_sprite(glm::vec2 pos, const sprite_t &sprite);
    void submit_quad(glm::vec2 pos, glm::vec2 scale);

    explicit Renderer(GLFWwindow *window);
    void create_pipeline_layout();
    void createFrameResources();
    void create_default_sampler();
    void recordCommandBuffer(
        VkCommandBuffer cmd,
        VkImage image,
        VkImageView imageView,
        VkExtent2D extent,
        VkImageLayout oldLayout,
        const Frame &frame,
        const std::vector<vertex_p2_s2_st2_col4_a1_u32> &vertices) const;
    void immediate_submit(std::function<void(VkCommandBuffer)>&& func) const;
    void transition_image_layout(VkCommandBuffer cmd, VkImage image, VkImageLayout oldLayout, VkImageLayout newLayout) const;
    VkImageView create_image_view(VkImage image, VkFormat format) const;
    AllocatedBuffer create_buffer(size_t allocSize, VkBufferUsageFlags usage, VmaMemoryUsage memoryUsage);
    void destroy_buffer(const AllocatedBuffer& buffer);
    // GPUMeshBuffers uploadMesh(std::span<uint32_t> indices, std::span<vertex_p2_st2_col4_a1_u32> vertices);
    void upload_vertex_buffer(
        VkCommandBuffer cmd,
        const Frame &frame,
        std::span<const vertex_p2_s2_st2_col4_a1_u32> vertices) const;

    [[nodiscard]] VkPipeline get_pipeline(PipelineType type) const;
    // void bind_material(VkCommandBuffer cmd, uint16_t materialID);
    void create_descriptor_pool();
    void update_bindless_slot(uint32_t slot, VkImageView view, VkSampler sampler) const;

    // Returns the resource info so the Manager can store it
    void upload_texture(
        int w,
        int h,
        const void* pixels,
        VkImage *image,
        VmaAllocation *allocation,
        VkImageView *view,
        uint32_t *descriptor_index);

    template <typename T>
    VkPipeline create_graphics_pipeline(
        VkDevice device,
        VkPipelineLayout layout,
        VkFormat colorFormat,
        // VkShaderModule vertShader,
        // VkShaderModule fragShader,
        VkPrimitiveTopology topology,
        bool enableBlending)
    {

        auto slangTargets{ std::to_array<slang::TargetDesc>({ {
            .format = SLANG_SPIRV,
            .profile = slangGlobalSession->findProfile("spirv_1_4")
        } })};
        auto slangOptions{ std::to_array<slang::CompilerOptionEntry>({ {
            slang::CompilerOptionName::EmitSpirvDirectly,
            {slang::CompilerOptionValueKind::Int, 1}
        } })};
        slang::SessionDesc slangSessionDesc{
            .targets = slangTargets.data(),
            .targetCount = SlangInt(slangTargets.size()),
            .defaultMatrixLayoutMode = SLANG_MATRIX_LAYOUT_COLUMN_MAJOR,
            .compilerOptionEntries = slangOptions.data(),
            .compilerOptionEntryCount = uint32_t(slangOptions.size())
        };
        Slang::ComPtr<slang::ISession> slangSession;
        slangGlobalSession->createSession(slangSessionDesc, slangSession.writeRef());

        Slang::ComPtr<slang::IModule> slangModule{
            slangSession->loadModuleFromSource("triangle", "shaders/shader.slang", nullptr, nullptr)
        };
        Slang::ComPtr<ISlangBlob> spirv;
        slangModule->getTargetCode(0, spirv.writeRef());

        VkShaderModuleCreateInfo shaderModuleCI{
            .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
            .codeSize = spirv->getBufferSize(),
            .pCode = (uint32_t*)spirv->getBufferPointer()
        };
        VkShaderModule shaderModule{};
        vkCreateShaderModule(device, &shaderModuleCI, nullptr, &shaderModule);

        auto vsCode = loadFile("shaders/triangle.vert.spv");
        auto fsCode = loadFile("shaders/triangle.frag.spv");

        VkShaderModuleCreateInfo smci{
            .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO
        };

        smci.codeSize = vsCode.size();
        smci.pCode = reinterpret_cast<uint32_t*>(vsCode.data());
        // VkShaderModule vs;
        // vkCreateShaderModule(device, &smci, nullptr, &vs);

        smci.codeSize = fsCode.size();
        smci.pCode = reinterpret_cast<uint32_t*>(fsCode.data());
        // VkShaderModule fs;
        // vkCreateShaderModule(device, &smci, nullptr, &fs);

        // --- Shaders ---
        std::vector<VkPipelineShaderStageCreateInfo> shaderStages{
        { .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
          .stage = VK_SHADER_STAGE_VERTEX_BIT,
          .module = shaderModule, .pName = "main"},
        { .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
          .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
          .module = shaderModule, .pName = "main" }
        };

        // --- Vertex Input (Generic) ---
        auto binding = T::getBindingDescription();
        auto attrs = T::getAttributeDescriptions();

        // --- Vertex Input (Matching our vertex_p2_st2_col4 struct) ---
        VkPipelineVertexInputStateCreateInfo vi{
            .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
            .vertexBindingDescriptionCount = 1,
            .pVertexBindingDescriptions = &binding,
            .vertexAttributeDescriptionCount = attrs.size(),
            .pVertexAttributeDescriptions = attrs.data(),
        };

        // --- Input Assembly (Changes based on Topology parameter) ---
        VkPipelineInputAssemblyStateCreateInfo ia{VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO};
        ia.topology = topology;

        // --- Blending (Changes based on enableBlending parameter) ---
        VkPipelineColorBlendAttachmentState colorBlend{
            .blendEnable = enableBlending ? VK_TRUE : VK_FALSE,
            .srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA,
            .dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA,
            .colorBlendOp = VK_BLEND_OP_ADD,
            .srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE,
            .dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
            .alphaBlendOp = VK_BLEND_OP_ADD,
            .colorWriteMask = 0xF
        };

        // --- Boilerplate (Standard 2D Defaults) ---
        VkPipelineViewportStateCreateInfo vp{
            VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
            nullptr,
            0,
            1,
            nullptr,
            1,
            nullptr
        };

        VkPipelineRasterizationStateCreateInfo rs{
            VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
            nullptr,
            0,
            0,
            0,
            VK_POLYGON_MODE_FILL,
            VK_CULL_MODE_NONE,
            VK_FRONT_FACE_COUNTER_CLOCKWISE,
            0,
            0,
            0,
            0,
            1.0f
        };

        VkPipelineMultisampleStateCreateInfo ms{
            VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
            nullptr,
            0,
            VK_SAMPLE_COUNT_1_BIT
        };

        VkPipelineColorBlendStateCreateInfo cb{
            VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
            nullptr,
            0,
            0,
            VK_LOGIC_OP_AND,
            1,
            &colorBlend
        };

        VkDynamicState dyns[] = { VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR };
        VkPipelineDynamicStateCreateInfo ds{
            VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
            nullptr,
            0,
            2,
            dyns
        };

        VkPipelineRenderingCreateInfo rci{
            VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO,
            nullptr,
            0,
            1,
            &colorFormat
        };

        VkGraphicsPipelineCreateInfo gpci{
            .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
            .pNext = &rci,
            .stageCount = (uint32_t) shaderStages.size(),
            .pStages = shaderStages.data(),
            .pVertexInputState = &vi,
            .pInputAssemblyState = &ia,
            .pViewportState = &vp,
            .pRasterizationState = &rs,
            .pMultisampleState = &ms,
            .pColorBlendState = &cb,
            .pDynamicState = &ds,
            .layout = layout
        };

        VkPipeline pipeline;
        vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &gpci, nullptr, &pipeline);
        return pipeline;
    }
};

#endif //V_RENDERER_H