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

#ifndef V_RENDERER_H
#define V_RENDERER_H

#include <volk/volk.h>
#include <GLFW/glfw3.h>
#include <glm/glm.hpp>
#include <vma/vk_mem_alloc.h>
#include "sprite.h"
#include "texture.h"
#include <misc.h>

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

struct vertex_p2_col4 {
    glm::vec2 pos;
    glm::vec4 col;
};

struct vertex_p2_st2_col4 {
    glm::vec2 pos;
    glm::vec2 st;
    glm::vec4 col;
};

struct vertex_p2_st2_col4_a1_u32 {
    glm::vec2 pos;
    glm::vec2 st;
    glm::vec4 col;
    float     alpha;
    uint32_t  textureID;

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

    static std::vector<VkVertexInputAttributeDescription> getAttributeDescriptions() {
        return {
            {0, 0, VK_FORMAT_R32G32_SFLOAT, offsetof(vertex_p2_st2_col4_a1_u32, pos)},
            {1, 0, VK_FORMAT_R32G32_SFLOAT, offsetof(vertex_p2_st2_col4_a1_u32, st)},
            {2, 0, VK_FORMAT_R32G32B32A32_SFLOAT, offsetof(vertex_p2_st2_col4_a1_u32, col)},
            {3, 0, VK_FORMAT_R32_SFLOAT, offsetof(vertex_p2_st2_col4_a1_u32, alpha)},
            {4, 0, VK_FORMAT_R32_UINT, offsetof(vertex_p2_st2_col4_a1_u32, textureID)},
        };
    }
};

struct vertex_st2 {
    glm::vec2 st;
};

struct vertex_p2_scale2_rot1_st2 {
    glm::vec2 pos;
    glm::vec2 scale;
    float     rot;
    glm::vec2 st;
};

typedef vertex_st2 quad_st2[6];

typedef vertex_p2_scale2_rot1_st2 quad_pos2_scale2_rot1_st2[6];

// 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 position;
    glm::vec2 size;
    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 Renderer {
    std::vector<RenderCommand> commands;

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

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

    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;
    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;

    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;

    VkCommandPool commandPool;

    static constexpr uint32_t MAX_FRAMES_IN_FLIGHT = 2;
    static constexpr uint32_t MAX_VERTICES_PER_FRAME = 65536;

    struct Frame {
        VkCommandBuffer command_buffer{};
        VkSemaphore imageAvailable{};
        VkSemaphore renderFinished{};
    };

    Frame frames[MAX_FRAMES_IN_FLIGHT];
    uint32_t currentFrame = 0;

    VkSemaphore timelineSemaphore{};
    uint64_t frameValue = 0;

    VkBuffer vertexBuffer;
    VmaAllocation vertexAllocation;
    VmaAllocationInfo vertexAllocInfo;

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

    VkPipeline get_pipeline(PipelineType type) const;
    void bind_material(VkCommandBuffer cmd, uint16_t materialID);
    void create_descriptor_pool();
    VkDescriptorSet create_texture_descriptor(VkImageView imageView, VkSampler sampler);
    void update_bindless_slot(uint32_t slot, VkImageView view, VkSampler sampler);

    // Returns the resource info so the Manager can store it
    Texture upload_texture(int w, int h, void* pixels);

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

        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 ---
        VkPipelineShaderStageCreateInfo stages[2] = {
            {
                VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
                nullptr,
                0,
                VK_SHADER_STAGE_VERTEX_BIT,
                vs,
                "main"
            },
            {
                VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
                nullptr,
                0,
                VK_SHADER_STAGE_FRAGMENT_BIT,
                fs,
                "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 = 0,
            .pVertexBindingDescriptions = nullptr,
            .vertexAttributeDescriptionCount = 0,
            .pVertexAttributeDescriptions = nullptr,
        };

        // --- 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_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 = 2,
            .pStages = stages,
            .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