#include "ui.h"

#include <stdlib.h>
#include <string.h>
#include <d3dcompiler.h>
#include <windows.h>

// 为 MinGW 提供 COM 助手宏（MSVC 已内联提供）
#ifndef _MSC_VER
#define ID3D11Device_CreateBuffer(p,a,b,c)           (p)->lpVtbl->CreateBuffer(p,a,b,c)
#define ID3D11Device_CreateVertexShader(p,a,b,c,d)   (p)->lpVtbl->CreateVertexShader(p,a,b,c,d)
#define ID3D11Device_CreatePixelShader(p,a,b,c,d)    (p)->lpVtbl->CreatePixelShader(p,a,b,c,d)
#define ID3D11Device_CreateInputLayout(p,a,b,c,d,e)  (p)->lpVtbl->CreateInputLayout(p,a,b,c,d,e)
#define ID3D11Device_CreateBlendState(p,a,b)         (p)->lpVtbl->CreateBlendState(p,a,b)
#define ID3D11Device_CreateRasterizerState(p,a,b)    (p)->lpVtbl->CreateRasterizerState(p,a,b)
#define ID3D11Device_CreateTexture2D(p,a,b,c)        (p)->lpVtbl->CreateTexture2D(p,a,b,c)
#define ID3D11Device_CreateShaderResourceView(p,a,b,c) (p)->lpVtbl->CreateShaderResourceView(p,a,b,c)
#define ID3D11Device_CreateSamplerState(p,a,b)       (p)->lpVtbl->CreateSamplerState(p,a,b)

#define ID3D11RasterizerState_Release(p)             (p)->lpVtbl->Release(p)
#define ID3D11BlendState_Release(p)                  (p)->lpVtbl->Release(p)
#define ID3D11InputLayout_Release(p)                 (p)->lpVtbl->Release(p)
#define ID3D11PixelShader_Release(p)                 (p)->lpVtbl->Release(p)
#define ID3D11VertexShader_Release(p)                (p)->lpVtbl->Release(p)
#define ID3D11Buffer_Release(p)                      (p)->lpVtbl->Release(p)
#define ID3D11Texture2D_Release(p)                   (p)->lpVtbl->Release(p)
#define ID3D11ShaderResourceView_Release(p)          (p)->lpVtbl->Release(p)
#define ID3D11SamplerState_Release(p)                (p)->lpVtbl->Release(p)
#define ID3D10Blob_Release(p)                        (p)->lpVtbl->Release(p)
#define ID3D10Blob_GetBufferPointer(p)               (p)->lpVtbl->GetBufferPointer(p)
#define ID3D10Blob_GetBufferSize(p)                  (p)->lpVtbl->GetBufferSize(p)

#define ID3D11DeviceContext_Map(p,a,b,c,d,e)         (p)->lpVtbl->Map(p,a,b,c,d,e)
#define ID3D11DeviceContext_Unmap(p,a,b)             (p)->lpVtbl->Unmap(p,a,b)
#define ID3D11DeviceContext_IASetVertexBuffers(p,a,b,c,d,e) (p)->lpVtbl->IASetVertexBuffers(p,a,b,c,d,e)
#define ID3D11DeviceContext_IASetIndexBuffer(p,a,b,c) (p)->lpVtbl->IASetIndexBuffer(p,a,b,c)
#define ID3D11DeviceContext_IASetPrimitiveTopology(p,a) (p)->lpVtbl->IASetPrimitiveTopology(p,a)
#define ID3D11DeviceContext_IASetInputLayout(p,a)    (p)->lpVtbl->IASetInputLayout(p,a)
#define ID3D11DeviceContext_VSSetShader(p,a,b,c)     (p)->lpVtbl->VSSetShader(p,a,b,c)
#define ID3D11DeviceContext_PSSetShader(p,a,b,c)     (p)->lpVtbl->PSSetShader(p,a,b,c)
#define ID3D11DeviceContext_PSSetShaderResources(p,a,b,c) (p)->lpVtbl->PSSetShaderResources(p,a,b,c)
#define ID3D11DeviceContext_PSSetSamplers(p,a,b,c)   (p)->lpVtbl->PSSetSamplers(p,a,b,c)
#define ID3D11DeviceContext_OMSetBlendState(p,a,b,c) (p)->lpVtbl->OMSetBlendState(p,a,b,c)
#define ID3D11DeviceContext_RSSetState(p,a)          (p)->lpVtbl->RSSetState(p,a)
#define ID3D11DeviceContext_DrawIndexed(p,a,b,c)     (p)->lpVtbl->DrawIndexed(p,a,b,c)
#endif

static const char* ui_vs = R"(
struct VS_INPUT { float2 Pos : POSITION; float2 UV : TEXCOORD0; uint Color : COLOR0; };
struct VS_OUTPUT { float4 Pos : SV_POSITION; float2 UV : TEXCOORD0; float4 Color : COLOR0; };
VS_OUTPUT main(VS_INPUT input) {
    VS_OUTPUT output;
    output.Pos = float4(input.Pos.x, input.Pos.y, 0, 1);
    output.UV = input.UV;
    output.Color = float4((input.Color & 0xFF) / 255.0, ((input.Color >> 8) & 0xFF) / 255.0,
                          ((input.Color >> 16) & 0xFF) / 255.0, ((input.Color >> 24) & 0xFF) / 255.0);
    return output;
}
)";

static const char* ui_ps = R"(
struct PS_INPUT { float4 Pos : SV_POSITION; float2 UV : TEXCOORD0; float4 Color : COLOR0; };
Texture2D tex : register(t0);
SamplerState samp : register(s0);
float4 main(PS_INPUT input) : SV_TARGET { return input.Color * tex.Sample(samp, input.UV); }
)";

#define UI_FONT_COLS 16
#define UI_FONT_ROWS 8
#define UI_FONT_CELL_W 16
#define UI_FONT_CELL_H 16
#define UI_FONT_ATLAS_W (UI_FONT_COLS * UI_FONT_CELL_W)
#define UI_FONT_ATLAS_H (UI_FONT_ROWS * UI_FONT_CELL_H)

static int ui_create_font_texture(UI* ui, ID3D11Device* device) {
    BITMAPINFO bmi = {0};
    bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bmi.bmiHeader.biWidth = UI_FONT_ATLAS_W;
    bmi.bmiHeader.biHeight = -UI_FONT_ATLAS_H;
    bmi.bmiHeader.biPlanes = 1;
    bmi.bmiHeader.biBitCount = 32;
    bmi.bmiHeader.biCompression = BI_RGB;

    HDC hdc = GetDC(NULL);
    HDC memdc = CreateCompatibleDC(hdc);
    void* bits = NULL;
    HBITMAP bmp = CreateDIBSection(memdc, &bmi, DIB_RGB_COLORS, &bits, NULL, 0);
    if (!bmp) {
        DeleteDC(memdc);
        ReleaseDC(NULL, hdc);
        return 0;
    }
    SelectObject(memdc, bmp);
    PatBlt(memdc, 0, 0, UI_FONT_ATLAS_W, UI_FONT_ATLAS_H, BLACKNESS);
    SetBkColor(memdc, RGB(0, 0, 0));
    SetTextColor(memdc, RGB(255, 255, 255));

    HFONT font = CreateFont(UI_FONT_CELL_H - 2, 0, 0, 0, FW_NORMAL, FALSE, FALSE, FALSE,
                            DEFAULT_CHARSET, OUT_OUTLINE_PRECIS, CLIP_DEFAULT_PRECIS,
                            CLEARTYPE_QUALITY, FIXED_PITCH | FF_MODERN, "Consolas");
    if (font) SelectObject(memdc, font);

    for (int i = 32; i < 128; i++) {
        int col = i % UI_FONT_COLS;
        int row = i / UI_FONT_COLS;
        char c = (char)i;
        TextOutA(memdc, col * UI_FONT_CELL_W + 1, row * UI_FONT_CELL_H + 1, &c, 1);

        ui->char_uvs[i][0] = (float)(col * UI_FONT_CELL_W) / UI_FONT_ATLAS_W;
        ui->char_uvs[i][1] = (float)(row * UI_FONT_CELL_H) / UI_FONT_ATLAS_H;
        ui->char_uvs[i][2] = (float)((col + 1) * UI_FONT_CELL_W) / UI_FONT_ATLAS_W;
        ui->char_uvs[i][3] = (float)((row + 1) * UI_FONT_CELL_H) / UI_FONT_ATLAS_H;
    }

    // 将亮度转为 alpha，文字存储为白色
    uint32_t* pixels = (uint32_t*)bits;
    for (int i = 0; i < UI_FONT_ATLAS_W * UI_FONT_ATLAS_H; i++) {
        uint32_t p = pixels[i];
        uint8_t b = (uint8_t)(p & 0xFF);
        uint8_t g = (uint8_t)((p >> 8) & 0xFF);
        uint8_t r = (uint8_t)((p >> 16) & 0xFF);
        uint8_t lum = (uint8_t)((r * 77 + g * 150 + b * 29) >> 8);
        pixels[i] = 0x00FFFFFF | ((uint32_t)lum << 24);
    }
    // 左上角像素保留为纯白，供纯色矩形采样
    pixels[0] = 0xFFFFFFFF;

    D3D11_TEXTURE2D_DESC desc = {0};
    desc.Width = UI_FONT_ATLAS_W;
    desc.Height = UI_FONT_ATLAS_H;
    desc.MipLevels = 1;
    desc.ArraySize = 1;
    desc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
    desc.SampleDesc.Count = 1;
    desc.Usage = D3D11_USAGE_IMMUTABLE;
    desc.BindFlags = D3D11_BIND_SHADER_RESOURCE;

    D3D11_SUBRESOURCE_DATA init = {0};
    init.pSysMem = bits;
    init.SysMemPitch = UI_FONT_ATLAS_W * 4;

    ID3D11Texture2D* tex = NULL;
    HRESULT hr = ID3D11Device_CreateTexture2D(device, &desc, &init, &tex);
    if (SUCCEEDED(hr) && tex) {
        ID3D11Device_CreateShaderResourceView(device, (ID3D11Resource*)tex, NULL, &ui->font_srv);
        ID3D11Texture2D_Release(tex);
    }

    D3D11_SAMPLER_DESC sd = {0};
    sd.Filter = D3D11_FILTER_MIN_MAG_MIP_POINT;
    sd.AddressU = D3D11_TEXTURE_ADDRESS_WRAP;
    sd.AddressV = D3D11_TEXTURE_ADDRESS_WRAP;
    sd.AddressW = D3D11_TEXTURE_ADDRESS_WRAP;
    ID3D11Device_CreateSamplerState(device, &sd, &ui->font_sampler);

    if (font) DeleteObject(font);
    DeleteObject(bmp);
    DeleteDC(memdc);
    ReleaseDC(NULL, hdc);

    return ui->font_srv != NULL && ui->font_sampler != NULL;
}

UI* ui_create(ID3D11Device* device) {
    UI* ui = (UI*)calloc(1, sizeof(UI));
    ui->font_size = 14.0f;
    ui->text_color = 0xFFFFFFFF;
    ui->bg_color = 0xFF333333;
    ui->border_color = 0xFF666666;
    ui->padding = 4.0f;

    D3D11_BUFFER_DESC vbd = {0};
    vbd.ByteWidth = sizeof(UIVertex) * UI_MAX_VERTICES;
    vbd.Usage = D3D11_USAGE_DYNAMIC;
    vbd.BindFlags = D3D11_BIND_VERTEX_BUFFER;
    vbd.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    ID3D11Device_CreateBuffer(device, &vbd, NULL, &ui->vertex_buffer);

    D3D11_BUFFER_DESC ibd = {0};
    ibd.ByteWidth = sizeof(uint32_t) * UI_MAX_INDICES;
    ibd.Usage = D3D11_USAGE_DYNAMIC;
    ibd.BindFlags = D3D11_BIND_INDEX_BUFFER;
    ibd.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    ID3D11Device_CreateBuffer(device, &ibd, NULL, &ui->index_buffer);

    // 编译着色器
    ID3DBlob* vs_blob = NULL, *ps_blob = NULL, *err = NULL;
    D3DCompile(ui_vs, strlen(ui_vs), NULL, NULL, NULL, "main", "vs_5_0", 0, 0, &vs_blob, &err);
    D3DCompile(ui_ps, strlen(ui_ps), NULL, NULL, NULL, "main", "ps_5_0", 0, 0, &ps_blob, &err);

    if (vs_blob && ps_blob) {
        ID3D11Device_CreateVertexShader(device, ID3D10Blob_GetBufferPointer(vs_blob),
            ID3D10Blob_GetBufferSize(vs_blob), NULL, &ui->vertex_shader);
        ID3D11Device_CreatePixelShader(device, ID3D10Blob_GetBufferPointer(ps_blob),
            ID3D10Blob_GetBufferSize(ps_blob), NULL, &ui->pixel_shader);

        D3D11_INPUT_ELEMENT_DESC layout[] = {
            {"POSITION", 0, DXGI_FORMAT_R32G32_FLOAT, 0, 0, D3D11_INPUT_PER_VERTEX_DATA, 0},
            {"TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT, 0, 8, D3D11_INPUT_PER_VERTEX_DATA, 0},
            {"COLOR", 0, DXGI_FORMAT_R8G8B8A8_UNORM, 0, 16, D3D11_INPUT_PER_VERTEX_DATA, 0},
        };
        ID3D11Device_CreateInputLayout(device, layout, 3,
            ID3D10Blob_GetBufferPointer(vs_blob), ID3D10Blob_GetBufferSize(vs_blob), &ui->input_layout);
        ID3D10Blob_Release(vs_blob);
        ID3D10Blob_Release(ps_blob);
    }

    D3D11_BLEND_DESC bd = {0};
    bd.RenderTarget[0].BlendEnable = TRUE;
    bd.RenderTarget[0].SrcBlend = D3D11_BLEND_SRC_ALPHA;
    bd.RenderTarget[0].DestBlend = D3D11_BLEND_INV_SRC_ALPHA;
    bd.RenderTarget[0].BlendOp = D3D11_BLEND_OP_ADD;
    bd.RenderTarget[0].SrcBlendAlpha = D3D11_BLEND_ONE;
    bd.RenderTarget[0].DestBlendAlpha = D3D11_BLEND_ZERO;
    bd.RenderTarget[0].BlendOpAlpha = D3D11_BLEND_OP_ADD;
    ID3D11Device_CreateBlendState(device, &bd, &ui->blend_state);

    D3D11_RASTERIZER_DESC rd = {0};
    rd.FillMode = D3D11_FILL_SOLID;
    rd.CullMode = D3D11_CULL_NONE;
    rd.ScissorEnable = TRUE;
    ID3D11Device_CreateRasterizerState(device, &rd, &ui->rasterizer);

    ui_create_font_texture(ui, device);

    return ui;
}

void ui_destroy(UI* ui) {
    if (!ui) return;
    if (ui->font_sampler) ID3D11SamplerState_Release(ui->font_sampler);
    if (ui->font_srv) ID3D11ShaderResourceView_Release(ui->font_srv);
    if (ui->rasterizer) ID3D11RasterizerState_Release(ui->rasterizer);
    if (ui->blend_state) ID3D11BlendState_Release(ui->blend_state);
    if (ui->input_layout) ID3D11InputLayout_Release(ui->input_layout);
    if (ui->pixel_shader) ID3D11PixelShader_Release(ui->pixel_shader);
    if (ui->vertex_shader) ID3D11VertexShader_Release(ui->vertex_shader);
    if (ui->index_buffer) ID3D11Buffer_Release(ui->index_buffer);
    if (ui->vertex_buffer) ID3D11Buffer_Release(ui->vertex_buffer);
    free(ui);
}

void ui_begin(UI* ui, float mx, float my, int md, int mp, int mr) {
    ui->mouse_x = mx;
    ui->mouse_y = my;
    ui->mouse_down = md;
    ui->mouse_pressed = mp;
    ui->mouse_released = mr;
    ui->vertex_count = 0;
    ui->index_count = 0;
}

void ui_end(UI* ui) {
    (void)ui;
}

void ui_draw_rect(UI* ui, float x, float y, float w, float h, uint32_t color) {
    uint32_t base = ui->vertex_count;
    // UV 指向字体图集左上角纯白像素，使纯色矩形不受纹理内容影响
    ui->vertices[ui->vertex_count++] = (UIVertex){x, y, 0.0f, 0.0f, color};
    ui->vertices[ui->vertex_count++] = (UIVertex){x + w, y, 0.0f, 0.0f, color};
    ui->vertices[ui->vertex_count++] = (UIVertex){x + w, y + h, 0.0f, 0.0f, color};
    ui->vertices[ui->vertex_count++] = (UIVertex){x, y + h, 0.0f, 0.0f, color};
    ui->indices[ui->index_count++] = base;
    ui->indices[ui->index_count++] = base + 1;
    ui->indices[ui->index_count++] = base + 2;
    ui->indices[ui->index_count++] = base;
    ui->indices[ui->index_count++] = base + 2;
    ui->indices[ui->index_count++] = base + 3;
}

void ui_draw_text(UI* ui, float x, float y, const char* text, uint32_t color) {
    if (!text) return;
    float px = x;
    float py = y;
    const float cw = (float)UI_FONT_CELL_W;
    const float ch = (float)UI_FONT_CELL_H;
    while (*text) {
        unsigned char c = (unsigned char)*text++;
        if (c < 32 || c >= 128) c = '?';
        float u1 = ui->char_uvs[c][0];
        float v1 = ui->char_uvs[c][1];
        float u2 = ui->char_uvs[c][2];
        float v2 = ui->char_uvs[c][3];

        uint32_t base = ui->vertex_count;
        ui->vertices[ui->vertex_count++] = (UIVertex){px, py, u1, v1, color};
        ui->vertices[ui->vertex_count++] = (UIVertex){px + cw, py, u2, v1, color};
        ui->vertices[ui->vertex_count++] = (UIVertex){px + cw, py + ch, u2, v2, color};
        ui->vertices[ui->vertex_count++] = (UIVertex){px, py + ch, u1, v2, color};
        ui->indices[ui->index_count++] = base;
        ui->indices[ui->index_count++] = base + 1;
        ui->indices[ui->index_count++] = base + 2;
        ui->indices[ui->index_count++] = base;
        ui->indices[ui->index_count++] = base + 2;
        ui->indices[ui->index_count++] = base + 3;
        px += cw;
    }
}

void ui_draw_rect_border(UI* ui, float x, float y, float w, float h,
                          uint32_t fill, uint32_t border, float thickness) {
    ui_draw_rect(ui, x, y, w, h, fill);
    ui_draw_rect(ui, x, y, w, thickness, border);
    ui_draw_rect(ui, x, y + h - thickness, w, thickness, border);
    ui_draw_rect(ui, x, y, thickness, h, border);
    ui_draw_rect(ui, x + w - thickness, y, thickness, h, border);
}

int ui_button(UI* ui, float x, float y, float w, float h, const char* text) {
    int hover = ui->mouse_x >= x && ui->mouse_x <= x + w &&
                ui->mouse_y >= y && ui->mouse_y <= y + h;
    uint32_t color = hover ? 0xFF555555 : 0xFF444444;
    if (hover && ui->mouse_pressed) color = 0xFF666666;
    ui_draw_rect_border(ui, x, y, w, h, color, 0xFF888888, 1);
    if (text) {
        float tx = x + ui->padding;
        float ty = y + (h - UI_FONT_CELL_H) * 0.5f;
        ui_draw_text(ui, tx, ty, text, ui->text_color);
    }
    return hover && ui->mouse_pressed;
}

float ui_slider(UI* ui, float x, float y, float w, float h, float value, float min_val, float max_val) {
    ui_draw_rect_border(ui, x, y, w, h, 0xFF333333, 0xFF666666, 1);

    int hover = ui->mouse_x >= x && ui->mouse_x <= x + w &&
                ui->mouse_y >= y && ui->mouse_y <= y + h;
    if (hover && ui->mouse_down) {
        float t = (ui->mouse_x - x) / w;
        value = min_val + t * (max_val - min_val);
        if (value < min_val) value = min_val;
        if (value > max_val) value = max_val;
    }

    float knob_x = x + ((value - min_val) / (max_val - min_val)) * (w - h);
    ui_draw_rect(ui, knob_x, y, h, h, 0xFF0078D4);
    return value;
}

int ui_checkbox(UI* ui, float x, float y, const char* text, int checked) {
    float s = 16.0f;
    ui_draw_rect_border(ui, x, y, s, s, checked ? 0xFF0078D4 : 0xFF333333, 0xFF666666, 1);
    if (checked) {
        // 绘制对勾
        ui_draw_rect(ui, x + 4, y + 7, 2, 2, 0xFFFFFFFF);
        ui_draw_rect(ui, x + 6, y + 9, 2, 2, 0xFFFFFFFF);
        ui_draw_rect(ui, x + 8, y + 5, 5, 2, 0xFFFFFFFF);
    }
    if (text) {
        ui_draw_text(ui, x + s + ui->padding, y + (s - UI_FONT_CELL_H) * 0.5f, text, ui->text_color);
    }
    if (ui->mouse_x >= x && ui->mouse_x <= x + s &&
        ui->mouse_y >= y && ui->mouse_y <= y + s && ui->mouse_pressed) {
        return !checked;
    }
    return checked;
}

void ui_label(UI* ui, float x, float y, const char* text) {
    if (text) {
        ui_draw_text(ui, x, y, text, ui->text_color);
    }
}

void ui_render(UI* ui, ID3D11DeviceContext* ctx) {
    if (ui->vertex_count == 0) return;

    D3D11_MAPPED_SUBRESOURCE mapped;
    if (SUCCEEDED(ID3D11DeviceContext_Map(ctx, (ID3D11Resource*)ui->vertex_buffer, 0,
                                          D3D11_MAP_WRITE_DISCARD, 0, &mapped))) {
        memcpy(mapped.pData, ui->vertices, sizeof(UIVertex) * ui->vertex_count);
        ID3D11DeviceContext_Unmap(ctx, (ID3D11Resource*)ui->vertex_buffer, 0);
    }
    if (SUCCEEDED(ID3D11DeviceContext_Map(ctx, (ID3D11Resource*)ui->index_buffer, 0,
                                          D3D11_MAP_WRITE_DISCARD, 0, &mapped))) {
        memcpy(mapped.pData, ui->indices, sizeof(uint32_t) * ui->index_count);
        ID3D11DeviceContext_Unmap(ctx, (ID3D11Resource*)ui->index_buffer, 0);
    }

    UINT stride = sizeof(UIVertex);
    UINT offset = 0;
    ID3D11DeviceContext_IASetVertexBuffers(ctx, 0, 1, &ui->vertex_buffer, &stride, &offset);
    ID3D11DeviceContext_IASetIndexBuffer(ctx, ui->index_buffer, DXGI_FORMAT_R32_UINT, 0);
    ID3D11DeviceContext_IASetPrimitiveTopology(ctx, D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
    ID3D11DeviceContext_IASetInputLayout(ctx, ui->input_layout);

    ID3D11DeviceContext_VSSetShader(ctx, ui->vertex_shader, NULL, 0);
    ID3D11DeviceContext_PSSetShader(ctx, ui->pixel_shader, NULL, 0);
    ID3D11DeviceContext_PSSetShaderResources(ctx, 0, 1, &ui->font_srv);
    ID3D11DeviceContext_PSSetSamplers(ctx, 0, 1, &ui->font_sampler);

    float blend[4] = {1, 1, 1, 1};
    ID3D11DeviceContext_OMSetBlendState(ctx, ui->blend_state, blend, 0xFFFFFFFF);
    ID3D11DeviceContext_RSSetState(ctx, ui->rasterizer);

    ID3D11DeviceContext_DrawIndexed(ctx, ui->index_count, 0, 0);

    ID3D11DeviceContext_OMSetBlendState(ctx, NULL, blend, 0xFFFFFFFF);
    ID3D11DeviceContext_RSSetState(ctx, NULL);
}