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

#include <print>
#include <slang.h>
#include <slang-com-ptr.h>
#include "../graphics.h"
#include "../sprite.h"
#include "renderer.h"

#include "metal.h"
#include "vertex_data.h"

extern int32_t window_width;
extern int32_t window_height;

extern Device                    metal_device;
extern MTL::CommandQueue        *queue;
extern CA::MetalLayer           *metal_layer;
extern CA::MetalDrawable        *metal_drawable;

bool SortKey::operator<(const SortKey& b) const {
    if (depth != b.depth) return depth < b.depth;
    if (pipeline != b.pipeline) return pipeline < b.pipeline;
    return materialID < b.materialID;
}

Renderer::Renderer(GLFWwindow *window) {

    create_render_pipeline();

    frame_semaphore = dispatch_semaphore_create(kMaxFramesInFlight);
    current_frame = 0;

    for (Frame &frame : frames) {
        frame.vertex_buffer = metal_device.device->newBuffer(
            4 * 1024 * 1024,
            MTL::ResourceStorageModeShared
            );

        frame.uniform_buffer = metal_device.device->newBuffer(
            sizeof(simd::float4x4),
            MTL::ResourceStorageModeShared
            );
    }

    // colored_quad_pipeline = create_graphics_pipeline<vertex_p2_s2_st2_col4_a1_u32>(
    //     device,
    //     pipelineLayout,
    //     swapchain_format.format,
    //     VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
    //     true
    //     );
    // create_default_sampler();
    // create_descriptor_pool();
    // createFrameResources();
}

void Renderer::begin_frame() {
    commands.clear();
}

void Renderer::submit_quad(simd::float2 pos, simd::float2 scale) {
    RenderCommand cmd {};
    cmd.pipeline = PipelineType::ColoredQuad;
    cmd.key = {
        (uint16_t) pos.y,
        0,
        (uint8_t) PipelineType::ColoredQuad
    };

    cmd.colored_quad = {
        .pos = pos,
        .scale = scale,
        .colour = {0, 1, 1, 1},
    };

    commands.push_back(cmd);
}

void Renderer::submit_sprite(simd::float2 pos, const sprite_t &sprite) {
    RenderCommand cmd {};
    cmd.pipeline = PipelineType::TexturedQuad;
    cmd.key = {
        (uint16_t) pos.y,
        0,
        (uint8_t) PipelineType::TexturedQuad
    };

    const Texture &texture = texture_manager.textures[sprite.texture];

    cmd.textured_quad = {
        .pos = pos,
        .scale = { sprite.scale.x, sprite.scale.y },
        .uv0 = {0, 0},
        .uv1 = {1, 1},
        .colour = {1, 1, 1, 1},
        .texture = texture.p_texture->texture,
    };

    commands.push_back(cmd);

    // assert(started == true, "You can't submit without having started the renderer first.");
    // renderable : Renderable;
    // renderable.type = .Sprite;
    //
    // if sprite.window_space
    //     renderable.projection_type = .ORTHOGRAPHIC_WINDOW;
    // else
    //     renderable.projection_type = .ORTHOGRAPHIC_WORLD;
    //
    // renderable.pos                  = pos;
    // renderable.sprite.texture_sheet = sprite.texture_sheet;
    // renderable.sprite.texture_cell  = sprite.texture_cell;
    // renderable.sprite.origin        = sprite.origin;
    // renderable.sprite.scale         = sprite.scale;
    // renderable.sprite.colour        = sprite.colour;
    // renderable.sprite.alpha         = alpha;
    //
    // array_add(*renderer.renderable_list, renderable);
}

MTL::RenderPipelineState *Renderer::get_pipeline(PipelineType type) const {
    switch (type) {
    case PipelineType::TexturedQuad: return textured_quad_pipeline;
    case PipelineType::ColoredQuad:  return colored_quad_pipeline;
    case PipelineType::Line:         return line_pipeline;
    default:                         return {};
    }
}

// void Renderer::bind_material(VkCommandBuffer cmd, uint16_t materialID) {
//     // In a real app, you'd have an array/map: std::vector<VkDescriptorSet> textureSets;
//     VkDescriptorSet set = textureSets[materialID];
//
//     vkCmdBindDescriptorSets(
//         cmd,
//         VK_PIPELINE_BIND_POINT_GRAPHICS,
//         pipelineLayout, // Our shared layout
//         0,              // Starting at Set 0
//         1,              // Binding 1 set
//         &set,
//         0, nullptr
//     );
// }

void Renderer::end_frame(GLFWwindow *window) {
    send_render_command(window);
}

void Renderer::create_render_pipeline() {
    MTL::Function *vertex_shader{};
    MTL::Function *fragment_shader{};
    // load_metal_shader(
    //     "shaders/shaders.metallib",
    //     "vertex_main",
    //     "fragment_main",
    //     &vertex_shader,
    //     &fragment_shader);

    // 3. Load the Module
    Slang::ComPtr<slang::IBlob> diagnostics;
    Slang::ComPtr<slang::IModule> module(slangSession->loadModule("shader", diagnostics.writeRef()));

    if(diagnostics) {
        fprintf(stderr, "%s\n", (const char*) diagnostics->getBufferPointer());
    }

    Slang::ComPtr<slang::IEntryPoint> vertEntry;
    Slang::ComPtr<slang::IEntryPoint> fragEntry;
    module->findEntryPointByName("vs_main", vertEntry.writeRef());
    module->findEntryPointByName("fs_main", fragEntry.writeRef());

    slang::IComponentType* components[] = { module, vertEntry, fragEntry };
    Slang::ComPtr<slang::IComponentType> program;
    slangSession->createCompositeComponentType(components, 3, program.writeRef());

    Slang::ComPtr<slang::IBlob> code;
    Slang::ComPtr<slang::IBlob> diagnosticBlob;

    program->getTargetCode(
        0,
        code.writeRef(),
        diagnostics.writeRef()
        );

    if (diagnostics)
        std::println("{}", (const char*)diagnostics->getBufferPointer());

    std::println("Generated MSL:\n{}", (const char*)code->getBufferPointer());

    NS::Error *error = nullptr;
    MTL::Library *library = nullptr;
    NS::String* source = NS::String::string((const char *) code->getBufferPointer(), NS::UTF8StringEncoding);
    MTL::CompileOptions* opts = MTL::CompileOptions::alloc()->init();
    library = metal_device.device->newLibrary(source, opts, &error);
    opts->release();

    if (!library) {
        if (error) {
            std::println("Metal library compile error:");
            std::println("Domain: {}",
                error->domain()->utf8String());
            std::println("Code: {}",
                error->code());
            std::println("Description:\n{}",
                error->localizedDescription()->utf8String());
        }
        std::abort();
    }

    NS::String *vname = NS::String::string("vs_main", NS::UTF8StringEncoding);
    NS::String *fname = NS::String::string("fs_main", NS::UTF8StringEncoding);
    vertex_shader = library->newFunction(vname);
    fragment_shader = library->newFunction(fname);

    if (!vertex_shader) {
        std::println("vs_main not found in generated MSL");
        std::abort();
    }
    if (!fragment_shader) {
        std::println("fs_main not found in generated MSL");
        std::abort();
    }

    library->release();

    MTL::VertexDescriptor *vd = vertex_p2_s2_uv2_c4_a1::vertexDescriptor();

    MTL::RenderPipelineDescriptor* renderPipelineDescriptor = MTL::RenderPipelineDescriptor::alloc()->init();
    renderPipelineDescriptor->setLabel(NS::String::string("Triangle Rendering Pipeline", NS::ASCIIStringEncoding));
    renderPipelineDescriptor->setVertexFunction(vertex_shader);
    renderPipelineDescriptor->setFragmentFunction(fragment_shader);
    renderPipelineDescriptor->setVertexDescriptor(vd);
    assert(renderPipelineDescriptor);
    const MTL::PixelFormat pixel_format = metal_layer->pixelFormat();
    renderPipelineDescriptor->colorAttachments()->object(0)->setPixelFormat(pixel_format);

    textured_quad_pipeline = metal_device.device->newRenderPipelineState(renderPipelineDescriptor, &error);
    if (!textured_quad_pipeline) {
        if (error) {
            std::println("Pipeline error: {}",
                error->localizedDescription()->utf8String());
        }
        std::abort();
    }
    renderPipelineDescriptor->release();
    vd->release();
}

float4x4 make_ortho(float left, float right, float bottom, float top, float near, float far) {
    float sx =  2.0f / (right - left);
    float sy =  2.0f / (top - bottom);
    float sz =  1.0f / (far - near);

    float tx = -(right + left) / (right - left);
    float ty = -(bottom + top) / (top - bottom);
    float tz = -near           / (far - near);

    return float4x4(
        float4{sx,   0.0f, 0.0f, 0.0f},
        float4{0.0f, sy,   0.0f, 0.0f},
        float4{0.0f, 0.0f, sz,   0.0f},
        float4{tx,   ty,   tz,   1.0f}
    );
}

void Renderer::encode_render_command(GLFWwindow *window, MTL::RenderCommandEncoder *render_command_encoder) {

    std::vector<vertex_p2_s2_uv2_c4_a1> vertices;

    for (auto& cmd : commands) {

        switch (cmd.pipeline) {
        case PipelineType::ColoredQuad: {
            const auto &q = cmd.colored_quad;

            vertex_p2_s2_uv2_c4_a1 vTL = { q.pos, q.scale, {}, q.colour };
            vertex_p2_s2_uv2_c4_a1 vTR = { q.pos, q.scale, {}, q.colour };
            vertex_p2_s2_uv2_c4_a1 vBL = { q.pos, q.scale, {}, q.colour };
            vertex_p2_s2_uv2_c4_a1 vBR = { q.pos, q.scale, {}, q.colour };

            vertices.push_back(vTL);
            vertices.push_back(vBL);
            vertices.push_back(vTR);

            vertices.push_back(vTR);
            vertices.push_back(vBL);
            vertices.push_back(vBR);

            break;
        }
        case PipelineType::TexturedQuad: {
            const auto &q = cmd.textured_quad;

            vertex_p2_s2_uv2_c4_a1 vTL = { q.pos, q.scale, q.uv0, q.colour, 1.0 };
            vertex_p2_s2_uv2_c4_a1 vTR = { q.pos, q.scale, {q.uv1.x, q.uv0.y}, q.colour, 1.0 };
            vertex_p2_s2_uv2_c4_a1 vBL = { q.pos, q.scale, {q.uv0.x, q.uv1.y}, q.colour, 1.0 };
            vertex_p2_s2_uv2_c4_a1 vBR = { q.pos, q.scale, q.uv1, q.colour, 1.0};

            vertices.push_back(vTL);
            vertices.push_back(vBL);
            vertices.push_back(vTR);

            vertices.push_back(vTR);
            vertices.push_back(vBL);
            vertices.push_back(vBR);

            break;
        }
        default:
            break;
        }
    }

    const Frame &frame = frames[current_frame];

    memcpy(frame.vertex_buffer->contents(), vertices.data(), vertices.size() * sizeof(vertex_p2_s2_uv2_c4_a1));

    MTL::SamplerDescriptor* sampler_desc = MTL::SamplerDescriptor::alloc()->init();

    sampler_desc->setMinFilter(MTL::SamplerMinMagFilterLinear);
    sampler_desc->setMagFilter(MTL::SamplerMinMagFilterLinear);
    MTL::SamplerState* sampler = metal_device.device->newSamplerState(sampler_desc);

    int width, height;
    glfwGetFramebufferSize(window, &width, &height);
    simd::float4x4 ortho = make_ortho(0.0, window_width, window_height, 0.0, 1.0, 0.0);
    auto tm = simd::transpose(ortho);
    memcpy(frame.uniform_buffer->contents(), &tm, sizeof(ortho));
    render_command_encoder->setRenderPipelineState(textured_quad_pipeline);
    render_command_encoder->setVertexBuffer(frame.vertex_buffer, 0, 0);
    render_command_encoder->setVertexBuffer(frame.uniform_buffer, 0, 1);
    render_command_encoder->setFragmentSamplerState(sampler, 0);
    MTL::PrimitiveType type_triangle = MTL::PrimitiveTypeTriangle;
    NS::UInteger vertexStart = 0;
    NS::UInteger vertexCount = vertices.size();
    const Texture &texture = texture_manager.textures["assets/boy.png"];
    render_command_encoder->setFragmentTexture(texture.p_texture->texture, 0);
    render_command_encoder->drawPrimitives(type_triangle, vertexStart, vertexCount);

    sampler_desc->release();
}

void Renderer::send_render_command(GLFWwindow *window) {
    dispatch_semaphore_wait(frame_semaphore, DISPATCH_TIME_FOREVER);

    command_buffer = queue->commandBuffer();

    dispatch_semaphore_t sem = frame_semaphore;
    command_buffer->addCompletedHandler([sem](MTL::CommandBuffer* pBuf) {
        dispatch_semaphore_signal(sem);
    });

    MTL::RenderPassDescriptor* renderPassDescriptor = MTL::RenderPassDescriptor::alloc()->init();
    MTL::RenderPassColorAttachmentDescriptor *cd = renderPassDescriptor->colorAttachments()->object(0);
    cd->setTexture(metal_drawable->texture());
    cd->setLoadAction(MTL::LoadActionClear);
    cd->setClearColor(MTL::ClearColor(
        100.0f / 255.0f,
        149.0f / 255.0f,
        237.0f / 255.0f,
        1.0
        ));
    cd->setStoreAction(MTL::StoreActionStore);

    MTL::RenderCommandEncoder* renderCommandEncoder = command_buffer->renderCommandEncoder(renderPassDescriptor);
    encode_render_command(window, renderCommandEncoder);
    renderCommandEncoder->endEncoding();

    command_buffer->presentDrawable(metal_drawable);
    command_buffer->commit();
    command_buffer->waitUntilCompleted();

    current_frame = (current_frame + 1) % kMaxFramesInFlight;

    renderPassDescriptor->release();
}