#include "ui.h"

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

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; };
float4 main(PS_INPUT input) : SV_TARGET { return input.Color; }
)";

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

    return ui;
}

void ui_destroy(UI* ui) {
    if (!ui) return;
    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;
    ui->vertices[ui->vertex_count++] = (UIVertex){x, y, 0, 0, color};
    ui->vertices[ui->vertex_count++] = (UIVertex){x + w, y, 1, 0, color};
    ui->vertices[ui->vertex_count++] = (UIVertex){x + w, y + h, 1, 1, color};
    ui->vertices[ui->vertex_count++] = (UIVertex){x, y + h, 0, 1, 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_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) {
    (void)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);
    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) {
    (void)text;
    float s = 16.0f;
    ui_draw_rect_border(ui, x, y, s, s, checked ? 0xFF0078D4 : 0xFF333333, 0xFF666666, 1);
    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) {
    (void)text;
    ui_draw_rect(ui, x, y, (float)(strlen(text) * 8), ui->font_size, 0x00000000);
}

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

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