#include "UIRenderer.h"

#include <iostream>

#define NK_IMPLEMENTATION
#include <nuklear.h>
#include <glm/gtc/matrix_transform.hpp>

namespace Artifact {

namespace Events {

struct DrawUI {};

}

static std::array<nk_keys, static_cast<size_t>(Key::Last) + 1> ArtifactToNuklear {};

namespace {

static struct _nkInit {
    _nkInit() {
        // Populate the Artifact-to-Nuklear key map.
#define X(nuklear, artifact) ArtifactToNuklear[static_cast<int>(Key::artifact)] = nuklear
        X(NK_KEY_CTRL, ControlLeft);
        X(NK_KEY_CTRL, ControlRight);
        
        X(NK_KEY_SHIFT, ShiftLeft);
        X(NK_KEY_SHIFT, ShiftRight);
        
        X(NK_KEY_ENTER, Enter);
        
        X(NK_KEY_BACKSPACE, DeleteBackward);
        X(NK_KEY_DEL, DeleteForward);
        
        X(NK_KEY_LEFT, ArrowLeft);
        X(NK_KEY_RIGHT, ArrowRight);
        X(NK_KEY_UP, ArrowUp);
        X(NK_KEY_DOWN, ArrowDown);
        
        X(NK_KEY_TAB, Tab);
#undef X
    }
} __nkInit;

}

void UIRenderer::init() {
    printf("UI: %p", graphics);
    
    graphics->window->listen<Events::InputBegin>([this](auto ev) {
        nk_input_begin(&ctx);
    });
    
    graphics->window->listen<Events::InputEnd>([this](auto ev) {
        nk_input_end(&ctx);
    });
    
    graphics->window->listen<Events::CursorPosEvent>([this](auto ev) {
        nk_input_motion(&ctx, ev.x, ev.y);
    });
    
    graphics->window->listen<Events::ScrollEvent>([this](auto ev) {
        nk_input_scroll(&ctx, nk_vec2(ev.dx, ev.dy));
    });
    
    graphics->window->listen<Events::MouseEvent>([this](auto ev) {
        nk_input_button(&ctx, ev.button == Events::MOUSE_BUTTON_RIGHT ? NK_BUTTON_RIGHT : NK_BUTTON_LEFT, (int)ev.x, (int)ev.y, ev.state);
    });
    
    graphics->window->listen<Events::KeyDownEvent>([this](auto ev) {
        nk_input_key(&ctx, ArtifactToNuklear[static_cast<int>(ev.key)], true);
    });
    
    graphics->window->listen<Events::KeyUpEvent>([this](auto ev) {
        nk_input_key(&ctx, ArtifactToNuklear[static_cast<int>(ev.key)], false);
    });
    
    graphics->window->listen<Events::CharInputEvent>([this](auto ev) {
        nk_input_char(&ctx, ev.codepoint);
    });
    
    nk_init_default(&ctx, nullptr);
    
    nk_font_atlas_init_default(&atlas);
    nk_font_atlas_begin(&atlas);
    
    nk_font* font;
    void* fontData = nullptr;
    
    // FIXME: Automatic path resolution
    std::string assetPath = std::string(getenv("HOME")) + "/eclipse-workspace/ArtifactEngine/assets";
    
    FILE* fontFile = fopen((assetPath + "/fonts/Arial.ttf").c_str(), "rb");
    if (fontFile) {
        fseek(fontFile, 0, SEEK_END);
        long fontSize = ftell(fontFile);
        fseek(fontFile, 0, SEEK_SET);
        fontData = malloc(fontSize);
        fread(fontData, 1, fontSize, fontFile);
        fclose(fontFile);
        font = nk_font_atlas_add_from_memory(&atlas, fontData, fontSize, 13.0f, nullptr);
    } else {
        font = nk_font_atlas_add_default(&atlas, 13.0f, nullptr);
    }

    int fontWidth, fontHeight;
    const void* fontImage = nk_font_atlas_bake(&atlas, &fontWidth, &fontHeight, NK_FONT_ATLAS_RGBA32);
    fontTexture = graphics->createTextureFromData(fontImage, fontWidth, fontHeight);
    fontTextureView = wgpuTextureCreateView(fontTexture, nullptr);
    nk_font_atlas_end(&atlas, nk_handle_id(0), nullptr);
    nk_style_set_font(&ctx, &font->handle);
    if (fontData) free(fontData);

    unsigned char whitePixel[4] = {255, 255, 255, 255};
    dummyTexture = graphics->createTextureFromData(whitePixel, 1, 1);
    dummyTextureView = wgpuTextureCreateView(dummyTexture, nullptr);
    
    nk_buffer_init_default(&vertexBufferNK);
    nk_buffer_init_default(&indexBufferNK);
    nk_buffer_init_default(&commandBufferNK);
}

struct NKVertex {
    float pos[2];
    float uv[2];
    uint8_t color[4];
};

void UIRenderer::reload() {
    WGPUBufferDescriptor uiVertexBufferDesc = {};
    uiVertexBufferDesc.size = 1024 * 1024 * sizeof(NKVertex);
    uiVertexBufferDesc.usage = WGPUBufferUsage_Vertex | WGPUBufferUsage_CopyDst;
    uiVertexBufferDesc.label = WGPUStringView{"UI Vertex Buffer", WGPU_STRLEN};
    vertexBuffer = wgpuDeviceCreateBuffer(graphics->device, &uiVertexBufferDesc);

    WGPUBufferDescriptor uiIndexBufferDesc = {};
    uiIndexBufferDesc.size = 1024 * 1024 * sizeof(uint16_t);
    uiIndexBufferDesc.usage = WGPUBufferUsage_Index | WGPUBufferUsage_CopyDst;
    uiIndexBufferDesc.label = WGPUStringView{"UI Index Buffer", WGPU_STRLEN};
    indexBuffer = wgpuDeviceCreateBuffer(graphics->device, &uiIndexBufferDesc);

    WGPUBufferDescriptor uiUniformBufferDesc = {};
    uiUniformBufferDesc.size = sizeof(glm::mat4);
    uiUniformBufferDesc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
    uiUniformBufferDesc.label = WGPUStringView{"UI Uniform Buffer", WGPU_STRLEN};
    uniformBuffer = wgpuDeviceCreateBuffer(graphics->device, &uiUniformBufferDesc);

    std::string uiVSCode = R"(
        struct Uniforms {
            ortho: mat4x4<f32>,
        }
        @group(0) @binding(0) var<uniform> uniforms: Uniforms;
        struct VertexOutput {
            @builtin(position) pos: vec4<f32>,
            @location(0) uv: vec2<f32>,
            @location(1) color: vec4<f32>,
        }
        @vertex
        fn vs_main(@location(0) pos: vec2<f32>, @location(1) uv: vec2<f32>, @location(2) color: vec4<f32>) -> VertexOutput {
            var out: VertexOutput;
            out.pos = uniforms.ortho * vec4<f32>(pos, 0.0, 1.0);
            out.uv = uv;
            out.color = color;
            return out;
        }
    )";
    std::string uiFSCode = R"(
        @group(0) @binding(1) var texture: texture_2d<f32>;
        @group(0) @binding(2) var sampler_: sampler;
        @fragment
        fn fs_main(@location(0) uv: vec2<f32>, @location(1) color: vec4<f32>) -> @location(0) vec4<f32> {
            return textureSample(texture, sampler_, uv) * color;
        }
    )";

    WGPUShaderModuleDescriptor uiVSDesc = {};
    WGPUShaderSourceWGSL src = {.code = WGPUStringView{uiVSCode.c_str(), uiVSCode.size()}};
    src.chain.sType = WGPUSType_ShaderSourceWGSL;
    uiVSDesc.nextInChain = &src.chain;
    WGPUShaderModule uiVSModule = wgpuDeviceCreateShaderModule(graphics->device, &uiVSDesc);

    WGPUShaderModuleDescriptor uiFSDesc = {};
    src = {.code = WGPUStringView{uiFSCode.c_str(), uiFSCode.size()}};
    src.chain.sType = WGPUSType_ShaderSourceWGSL;
    uiFSDesc.nextInChain = &src.chain;
    WGPUShaderModule uiFSModule = wgpuDeviceCreateShaderModule(graphics->device, &uiFSDesc);

    WGPUSamplerDescriptor uiSamplerDesc = {};
    uiSamplerDesc.addressModeU = WGPUAddressMode_ClampToEdge;
    uiSamplerDesc.addressModeV = WGPUAddressMode_ClampToEdge;
    uiSamplerDesc.addressModeW = WGPUAddressMode_ClampToEdge;
    uiSamplerDesc.magFilter = WGPUFilterMode_Linear;
    uiSamplerDesc.minFilter = WGPUFilterMode_Linear;
    uiSamplerDesc.mipmapFilter = WGPUMipmapFilterMode_Linear;
    uiSamplerDesc.maxAnisotropy = 1;
    sampler = wgpuDeviceCreateSampler(graphics->device, &uiSamplerDesc);
    if (!sampler) {
        throw std::runtime_error("Failed to create UI sampler");
    }

    WGPUBindGroupLayoutEntry uiBglEntries[3] = {};
    uiBglEntries[0].binding = 0;
    uiBglEntries[0].visibility = WGPUShaderStage_Vertex;
    uiBglEntries[0].buffer.type = WGPUBufferBindingType_Uniform;
    uiBglEntries[1].binding = 1;
    uiBglEntries[1].visibility = WGPUShaderStage_Fragment;
    uiBglEntries[1].texture.sampleType = WGPUTextureSampleType_Float;
    uiBglEntries[1].texture.viewDimension = WGPUTextureViewDimension_2D;
    uiBglEntries[2].binding = 2;
    uiBglEntries[2].visibility = WGPUShaderStage_Fragment;
    uiBglEntries[2].sampler.type = WGPUSamplerBindingType_Filtering;

    WGPUBindGroupLayoutDescriptor uiBglDesc = {};
    uiBglDesc.entryCount = 3;
    uiBglDesc.entries = uiBglEntries;
    bgl = wgpuDeviceCreateBindGroupLayout(graphics->device, &uiBglDesc);
    if (!bgl) {
        throw std::runtime_error("Failed to create bind group layout");
    }

    WGPUVertexAttribute uiAttributes[3] = {};
    uiAttributes[0].format = WGPUVertexFormat_Float32x2;
    uiAttributes[0].offset = offsetof(NKVertex, pos);
    uiAttributes[0].shaderLocation = 0;
    uiAttributes[1].format = WGPUVertexFormat_Float32x2;
    uiAttributes[1].offset = offsetof(NKVertex, uv);
    uiAttributes[1].shaderLocation = 1;
    uiAttributes[2].format = WGPUVertexFormat_Unorm8x4;
    uiAttributes[2].offset = offsetof(NKVertex, color);
    uiAttributes[2].shaderLocation = 2;

    WGPUVertexBufferLayout uiVBLayout = {};
    uiVBLayout.arrayStride = sizeof(NKVertex);
    uiVBLayout.stepMode = WGPUVertexStepMode_Vertex;
    uiVBLayout.attributeCount = 3;
    uiVBLayout.attributes = uiAttributes;

    WGPUBlendState uiBlendState = {};
    uiBlendState.color.srcFactor = WGPUBlendFactor_SrcAlpha;
    uiBlendState.color.dstFactor = WGPUBlendFactor_OneMinusSrcAlpha;
    uiBlendState.color.operation = WGPUBlendOperation_Add;
    uiBlendState.alpha.srcFactor = WGPUBlendFactor_One;
    uiBlendState.alpha.dstFactor = WGPUBlendFactor_OneMinusSrcAlpha;
    uiBlendState.alpha.operation = WGPUBlendOperation_Add;

    WGPUColorTargetState uiColorTarget = {};
    uiColorTarget.format = WGPUTextureFormat_BGRA8Unorm;
    uiColorTarget.blend = &uiBlendState;
    uiColorTarget.writeMask = WGPUColorWriteMask_All;

    WGPUFragmentState uiFragmentState = {};
    uiFragmentState.module = uiFSModule;
    uiFragmentState.entryPoint = WGPUStringView{"fs_main", WGPU_STRLEN};
    uiFragmentState.targetCount = 1;
    uiFragmentState.targets = &uiColorTarget;

    WGPUPipelineLayoutDescriptor uiPlDesc = {};
    uiPlDesc.bindGroupLayoutCount = 1;
    uiPlDesc.bindGroupLayouts = &bgl;
    WGPUPipelineLayout uiPipelineLayout = wgpuDeviceCreatePipelineLayout(graphics->device, &uiPlDesc);

    WGPURenderPipelineDescriptor uiPipelineDesc = {};
    uiPipelineDesc.vertex.module = uiVSModule;
    uiPipelineDesc.vertex.entryPoint = WGPUStringView{"vs_main", WGPU_STRLEN};
    uiPipelineDesc.vertex.bufferCount = 1;
    uiPipelineDesc.vertex.buffers = &uiVBLayout;
    uiPipelineDesc.fragment = &uiFragmentState;
    uiPipelineDesc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
    uiPipelineDesc.primitive.frontFace = WGPUFrontFace_CCW;
    uiPipelineDesc.primitive.cullMode = WGPUCullMode_None;
    uiPipelineDesc.layout = uiPipelineLayout;
    uiPipelineDesc.multisample.count = 1;
    uiPipelineDesc.multisample.mask = ~0u;

    pipeline = wgpuDeviceCreateRenderPipeline(graphics->device, &uiPipelineDesc);
}

std::string txt = "Test";
int _txt = 0;

void UIRenderer::render(WGPUTextureView nextTexture, WGPUCommandEncoder encoder) {
    if (nk_begin(&ctx, "Test", nk_rect(10, 10, 100, 100), NK_WINDOW_TITLE | NK_WINDOW_BORDER | NK_WINDOW_MOVABLE | NK_WINDOW_SCALABLE)) {
        nk_layout_row_dynamic(&ctx, 50, 2);
        nk_label(&ctx, "Test", 0);
        nk_button_label(&ctx, "Test");
        
        nk_layout_row_dynamic(&ctx, 50, 2);
        nk_label(&ctx, "Test", 0);
        nk_button_label(&ctx, "Test");
        
        nk_layout_row_dynamic(&ctx, 50, 1);
        nk_edit_string(&ctx, NK_EDIT_BOX, txt.data(), &_txt, 100, nk_filter_ascii);
        
        if (nk_contextual_begin(&ctx, NK_PANEL_CONTEXTUAL, nk_vec2(100, 100), nk_rect(10, 10, 100, 100))) {
            nk_layout_row_dynamic(&ctx, 25, 1);
            nk_contextual_item_label(&ctx, "Button", NK_TEXT_CENTERED);
            nk_contextual_end(&ctx);
        }
    }
    
    nk_end(&ctx);
    
    dispatch(Events::DrawUI {});
    
    nk_buffer_clear(&vertexBufferNK);
    nk_buffer_clear(&indexBufferNK);
    nk_buffer_clear(&commandBufferNK);
    struct nk_convert_config config = {};
    static const struct nk_draw_vertex_layout_element vertex_layout[] = {
        {NK_VERTEX_POSITION, NK_FORMAT_FLOAT, NK_OFFSETOF(NKVertex, pos)},
        {NK_VERTEX_TEXCOORD, NK_FORMAT_FLOAT, NK_OFFSETOF(NKVertex, uv)},
        {NK_VERTEX_COLOR, NK_FORMAT_R8G8B8A8, NK_OFFSETOF(NKVertex, color)},
        {NK_VERTEX_LAYOUT_END}
    };
    config.vertex_layout = vertex_layout;
    config.vertex_size = sizeof(NKVertex);
    config.vertex_alignment = alignof(NKVertex);
    config.tex_null.texture = nk_handle_id(0);
    config.tex_null.uv = {0.0f, 0.0f};
    config.circle_segment_count = 22;
    config.curve_segment_count = 22;
    config.arc_segment_count = 22;
    config.global_alpha = 1.0f;
    config.shape_AA = NK_ANTI_ALIASING_ON;
    config.line_AA = NK_ANTI_ALIASING_ON;
    nk_convert(&ctx, &commandBufferNK, &vertexBufferNK, &indexBufferNK, &config);
    
    // FIXME: Variable size
    int viewportWidth = 1080;
    int viewportHeight = 640;
    
    auto ortho = glm::ortho(0.0f, (float)viewportWidth, (float)viewportHeight, 0.0f, -1.0f, 1.0f);
    wgpuQueueWriteBuffer(graphics->queue, uniformBuffer, 0, &ortho, sizeof(ortho));
    
    WGPURenderPassColorAttachment colorAttachment = {};
    colorAttachment.view = nextTexture;
    colorAttachment.loadOp = WGPULoadOp_Load;
    colorAttachment.storeOp = WGPUStoreOp_Store;
    colorAttachment.clearValue = {0, 0.2, 0.4, 1};
    colorAttachment.depthSlice = -1;

    WGPURenderPassDepthStencilAttachment depthAttachment = {};
    depthAttachment.view = graphics->depthTextureView;
    depthAttachment.depthLoadOp = WGPULoadOp_Clear;
    depthAttachment.depthStoreOp = WGPUStoreOp_Store;
    depthAttachment.depthClearValue = 1.0f;
    depthAttachment.depthReadOnly = false;
    depthAttachment.stencilLoadOp = WGPULoadOp_Undefined;
    depthAttachment.stencilStoreOp = WGPUStoreOp_Undefined;

    WGPURenderPassDescriptor renderPassDesc = {};
    renderPassDesc.colorAttachmentCount = 1;
    renderPassDesc.colorAttachments = &colorAttachment;
    renderPassDesc.depthStencilAttachment = &depthAttachment;

    colorAttachment.loadOp = WGPULoadOp_Load;
    WGPURenderPassDescriptor uiRenderPassDesc = {};
    uiRenderPassDesc.colorAttachmentCount = 1;
    uiRenderPassDesc.colorAttachments = &colorAttachment;

    WGPURenderPassEncoder uiRenderPass = wgpuCommandEncoderBeginRenderPass(encoder, &uiRenderPassDesc);
    wgpuRenderPassEncoderSetPipeline(uiRenderPass, pipeline);
    wgpuRenderPassEncoderSetViewport(uiRenderPass, 0, 0, (float)viewportWidth, (float)viewportHeight, 0.0f, 1.0f);


    size_t vertexSize = nk_buffer_total(&vertexBufferNK);
    size_t indexSize = nk_buffer_total(&indexBufferNK);
    if (vertexSize > 0 && indexSize > 0) {
        wgpuQueueWriteBuffer(graphics->queue, vertexBuffer, 0, vertexBufferNK.memory.ptr, vertexSize);
        wgpuQueueWriteBuffer(graphics->queue, indexBuffer, 0, indexBufferNK.memory.ptr, indexSize);

        wgpuRenderPassEncoderSetVertexBuffer(uiRenderPass, 0, vertexBuffer, 0, vertexSize);
        wgpuRenderPassEncoderSetIndexBuffer(uiRenderPass, indexBuffer, WGPUIndexFormat_Uint16, 0, indexSize);

        std::vector<WGPUBindGroup> bindGroups;
        WGPUBindGroupEntry uiBgEntries[3] = {};
        uiBgEntries[0].binding = 0;
        uiBgEntries[0].buffer = uniformBuffer;
        uiBgEntries[0].offset = 0;
        uiBgEntries[0].size = sizeof(glm::mat4);
        uiBgEntries[2].binding = 2;
        uiBgEntries[2].sampler = sampler;

        std::unordered_map<int, WGPUTextureView> textureMap = {{0, fontTextureView}};
        const struct nk_draw_command* cmd;
        uint32_t offset = 0;
        nk_draw_foreach(cmd, &ctx, &commandBufferNK) {
            if (!cmd->elem_count) continue;
            if (offset + cmd->elem_count > indexSize / sizeof(uint16_t)) break;

            uint32_t scissorX = static_cast<uint32_t>(cmd->clip_rect.x < 0 ? 0 : cmd->clip_rect.x);
            uint32_t scissorY = static_cast<uint32_t>(cmd->clip_rect.y < 0 ? 0 : cmd->clip_rect.y);
            uint32_t scissorW = static_cast<uint32_t>(cmd->clip_rect.w);
            uint32_t scissorH = static_cast<uint32_t>(cmd->clip_rect.h);
            scissorX = std::min(scissorX, (uint32_t)viewportWidth - 1);
            scissorY = std::min(scissorY, (uint32_t)viewportHeight - 1);
            if (scissorW > viewportWidth || scissorH > viewportHeight || scissorX + scissorW > viewportWidth || scissorY + scissorH > viewportHeight) {
                scissorW = std::min(scissorW, viewportWidth - scissorX);
                scissorH = std::min(scissorH, viewportHeight - scissorY);
            }
            wgpuRenderPassEncoderSetScissorRect(uiRenderPass, scissorX, scissorY, scissorW, scissorH);

            uiBgEntries[1].binding = 1;
            uiBgEntries[1].textureView = cmd->texture.id >= 0 && textureMap.count(cmd->texture.id) ? textureMap[cmd->texture.id] : dummyTextureView;

            WGPUBindGroupDescriptor uiBgDesc = {};
            uiBgDesc.layout = bgl;
            uiBgDesc.entryCount = 3;
            uiBgDesc.entries = uiBgEntries;

            WGPUBindGroup uiDynamicBindGroup = wgpuDeviceCreateBindGroup(graphics->device, &uiBgDesc);
            bindGroups.push_back(uiDynamicBindGroup);

            wgpuRenderPassEncoderSetBindGroup(uiRenderPass, 0, uiDynamicBindGroup, 0, nullptr);
            wgpuRenderPassEncoderDrawIndexed(uiRenderPass, cmd->elem_count, 1, offset, 0, 0);

            offset += cmd->elem_count;
        }
        wgpuRenderPassEncoderEnd(uiRenderPass);
        for (WGPUBindGroup bg : bindGroups) wgpuBindGroupRelease(bg);
    } else {
        wgpuRenderPassEncoderEnd(uiRenderPass);
    }

    nk_clear(&ctx);
}

}