#include "uilayer.h"
#include <stdlib.h>
#include <string.h>
#include "../Core/logger.h"

#define UI_MAX_VERTS (UI_MAX_QUADS * 4)
#define UI_MAX_IDXS (UI_MAX_QUADS * 6)

typedef struct { float p[3], n[3], uv[2], c[4]; } UIVert;

static UIVert g_verts[UI_MAX_VERTS];
static uint32_t g_indices[UI_MAX_IDXS];
static int g_vi = 0;
static int g_ii = 0;

static UIVert g_text_verts[UI_MAX_VERTS];
static uint32_t g_text_indices[UI_MAX_IDXS];
static int g_tvi = 0;
static int g_tii = 0;
static ID3D11ShaderResourceView* g_text_srv = NULL;

static void screen_to_world(float sx, float sy, float sw, float sh, float z, float* wx, float* wy) {
    *wx = sx - sw * 0.5f;
    *wy = -(sy - sh * 0.5f);
}

int uilayer_init(UILayer* u, ID3D11Device* dev, float screen_w, float screen_h) {
    memset(u, 0, sizeof(*u));
    u->screen_w = screen_w;
    u->screen_h = screen_h;
    u->z_plane = 0.0f;

    u->camera = camera_orthographic(screen_w, screen_h, 0.1f, 200.0f);
    u->camera.position = (Vec3){0, 0, 100};
    u->camera.target = (Vec3){0, 0, 0};
    u->camera.dirty = 1;
    camera_update(&u->camera);

    u->shader = shader_create();
    if (!shader_compile(u->shader, dev, NULL, NULL, 0, 0)) {
        LOG_ERROR("UILayer: shader compile failed");
        return 0;
    }

    D3D11_BUFFER_DESC bd = {0};
    D3D11_SUBRESOURCE_DATA sd = {0};

    UIVert* empty = (UIVert*)calloc(UI_MAX_VERTS, sizeof(UIVert));
    bd.ByteWidth = UI_MAX_VERTS * sizeof(UIVert);
    bd.Usage = D3D11_USAGE_DEFAULT;
    bd.BindFlags = D3D11_BIND_VERTEX_BUFFER;
    sd.pSysMem = empty;
    dev->lpVtbl->CreateBuffer(dev, &bd, &sd, &u->vb);
    free(empty);

    uint32_t* eidx = (uint32_t*)calloc(UI_MAX_IDXS, sizeof(uint32_t));
    bd.ByteWidth = UI_MAX_IDXS * sizeof(uint32_t);
    bd.BindFlags = D3D11_BIND_INDEX_BUFFER;
    sd.pSysMem = eidx;
    dev->lpVtbl->CreateBuffer(dev, &bd, &sd, &u->ib);
    free(eidx);

    UIVert* empty2 = (UIVert*)calloc(UI_MAX_VERTS, sizeof(UIVert));
    bd.ByteWidth = UI_MAX_VERTS * sizeof(UIVert);
    bd.Usage = D3D11_USAGE_DEFAULT;
    bd.BindFlags = D3D11_BIND_VERTEX_BUFFER;
    sd.pSysMem = empty2;
    dev->lpVtbl->CreateBuffer(dev, &bd, &sd, &u->text_vb);
    free(empty2);

    uint32_t* eidx2 = (uint32_t*)calloc(UI_MAX_IDXS, sizeof(uint32_t));
    bd.ByteWidth = UI_MAX_IDXS * sizeof(uint32_t);
    bd.BindFlags = D3D11_BIND_INDEX_BUFFER;
    sd.pSysMem = eidx2;
    dev->lpVtbl->CreateBuffer(dev, &bd, &sd, &u->text_ib);
    free(eidx2);

    D3D11_RASTERIZER_DESC rdesc = {0};
    rdesc.FillMode = D3D11_FILL_SOLID;
    rdesc.CullMode = D3D11_CULL_NONE;
    dev->lpVtbl->CreateRasterizerState(dev, &rdesc, &u->rstate);

    D3D11_DEPTH_STENCIL_DESC dsdesc = {0};
    dsdesc.DepthEnable = FALSE;
    dsdesc.DepthWriteMask = D3D11_DEPTH_WRITE_MASK_ZERO;
    dev->lpVtbl->CreateDepthStencilState(dev, &dsdesc, &u->ds_state);

    D3D11_BLEND_DESC bdesc = {0};
    bdesc.RenderTarget[0].BlendEnable = TRUE;
    bdesc.RenderTarget[0].SrcBlend = D3D11_BLEND_SRC_ALPHA;
    bdesc.RenderTarget[0].DestBlend = D3D11_BLEND_INV_SRC_ALPHA;
    bdesc.RenderTarget[0].BlendOp = D3D11_BLEND_OP_ADD;
    bdesc.RenderTarget[0].SrcBlendAlpha = D3D11_BLEND_ONE;
    bdesc.RenderTarget[0].DestBlendAlpha = D3D11_BLEND_ZERO;
    bdesc.RenderTarget[0].BlendOpAlpha = D3D11_BLEND_OP_ADD;
    dev->lpVtbl->CreateBlendState(dev, &bdesc, &u->text_blend);

    D3D11_SAMPLER_DESC sdesc = {0};
    sdesc.Filter = D3D11_FILTER_MIN_MAG_MIP_POINT;
    sdesc.AddressU = sdesc.AddressV = sdesc.AddressW = D3D11_TEXTURE_ADDRESS_CLAMP;
    dev->lpVtbl->CreateSamplerState(dev, &sdesc, &u->text_sampler);

    const char* text_vs_src =
        "cbuffer CB:register(b0){float4x4 WVP;};"
        "struct VS_IN{float4 pos:POSITION;float3 n:NORMAL;float2 uv:TEXCOORD;float4 col:COLOR;};"
        "struct PS_IN{float4 pos:SV_POSITION;float2 uv:TEXCOORD;float4 col:COLOR;};"
        "PS_IN main(VS_IN i){PS_IN o;o.pos=mul(i.pos,WVP);o.uv=i.uv;o.col=i.col;return o;}";
    const char* text_ps_src =
        "Texture2D tex:register(t0);SamplerState smp:register(s0);"
        "struct PS_IN{float4 pos:SV_POSITION;float2 uv:TEXCOORD;float4 col:COLOR;};"
        "float4 main(PS_IN i):SV_TARGET{float a=tex.Sample(smp,i.uv).a;return float4(i.col.rgb,i.col.a*a);}";

    ID3DBlob *vsb = NULL, *psb = NULL;
    D3DCompile(text_vs_src, (UINT)strlen(text_vs_src), NULL, NULL, NULL, "main", "vs_5_0", 0, 0, &vsb, NULL);
    D3DCompile(text_ps_src, (UINT)strlen(text_ps_src), NULL, NULL, NULL, "main", "ps_5_0", 0, 0, &psb, NULL);
    if (vsb && psb) {
        dev->lpVtbl->CreateVertexShader(dev, vsb->lpVtbl->GetBufferPointer(vsb), vsb->lpVtbl->GetBufferSize(vsb), NULL, &u->text_vs);
        dev->lpVtbl->CreatePixelShader(dev, psb->lpVtbl->GetBufferPointer(psb), psb->lpVtbl->GetBufferSize(psb), NULL, &u->text_ps);
        D3D11_INPUT_ELEMENT_DESC layout[] = {
            {"POSITION",0,DXGI_FORMAT_R32G32B32_FLOAT,0,0,D3D11_INPUT_PER_VERTEX_DATA,0},
            {"NORMAL",0,DXGI_FORMAT_R32G32B32_FLOAT,0,12,D3D11_INPUT_PER_VERTEX_DATA,0},
            {"TEXCOORD",0,DXGI_FORMAT_R32G32_FLOAT,0,24,D3D11_INPUT_PER_VERTEX_DATA,0},
            {"COLOR",0,DXGI_FORMAT_R32G32B32A32_FLOAT,0,32,D3D11_INPUT_PER_VERTEX_DATA,0},
        };
        dev->lpVtbl->CreateInputLayout(dev, layout, 4, vsb->lpVtbl->GetBufferPointer(vsb), vsb->lpVtbl->GetBufferSize(vsb), &u->text_layout);
        vsb->lpVtbl->Release(vsb);
        psb->lpVtbl->Release(psb);
    }

    LOG_INFO("UILayer: initialized (%.0fx%.0f)", screen_w, screen_h);
    return 1;
}

void uilayer_destroy(UILayer* u) {
    if (!u) return;
    if (u->text_sampler) u->text_sampler->lpVtbl->Release(u->text_sampler);
    if (u->text_blend) u->text_blend->lpVtbl->Release(u->text_blend);
    if (u->text_layout) u->text_layout->lpVtbl->Release(u->text_layout);
    if (u->text_ps) u->text_ps->lpVtbl->Release(u->text_ps);
    if (u->text_vs) u->text_vs->lpVtbl->Release(u->text_vs);
    if (u->text_ib) u->text_ib->lpVtbl->Release(u->text_ib);
    if (u->text_vb) u->text_vb->lpVtbl->Release(u->text_vb);
    if (u->ds_state) u->ds_state->lpVtbl->Release(u->ds_state);
    if (u->rstate) u->rstate->lpVtbl->Release(u->rstate);
    if (u->ib) u->ib->lpVtbl->Release(u->ib);
    if (u->vb) u->vb->lpVtbl->Release(u->vb);
    if (u->shader) shader_destroy(u->shader);
    memset(u, 0, sizeof(*u));
}

void uilayer_begin(UILayer* u) {
    (void)u;
    g_vi = 0; g_ii = 0;
    g_tvi = 0; g_tii = 0;
    g_text_srv = NULL;
}

void uilayer_rect(UILayer* u, float x, float y, float w, float h,
                   float r, float g, float b, float a) {
    if (g_vi + 4 > UI_MAX_VERTS) return;
    float x0, y0, x1, y1;
    screen_to_world(x, y, u->screen_w, u->screen_h, u->z_plane, &x0, &y0);
    screen_to_world(x + w, y + h, u->screen_w, u->screen_h, u->z_plane, &x1, &y1);
    float z = u->z_plane;
    uint32_t base = (uint32_t)g_vi;
    g_verts[g_vi] = (UIVert){{x0,y0,z},{0,0,1},{0,0},{r,g,b,a}}; g_vi++;
    g_verts[g_vi] = (UIVert){{x1,y0,z},{0,0,1},{1,0},{r,g,b,a}}; g_vi++;
    g_verts[g_vi] = (UIVert){{x1,y1,z},{0,0,1},{1,1},{r,g,b,a}}; g_vi++;
    g_verts[g_vi] = (UIVert){{x0,y1,z},{0,0,1},{0,1},{r,g,b,a}}; g_vi++;
    g_indices[g_ii++] = base; g_indices[g_ii++] = base+1; g_indices[g_ii++] = base+2;
    g_indices[g_ii++] = base; g_indices[g_ii++] = base+2; g_indices[g_ii++] = base+3;
}

void uilayer_textured_rect(UILayer* u, ID3D11DeviceContext* ctx,
                            ID3D11ShaderResourceView* srv,
                            float x, float y, float w, float h,
                            float u0, float v0, float u1, float v1,
                            float r, float g, float b, float a) {
    (void)ctx;
    if (!srv || g_tvi + 4 > UI_MAX_VERTS) return;
    float x0, y0, x1, y1;
    screen_to_world(x, y, u->screen_w, u->screen_h, u->z_plane, &x0, &y0);
    screen_to_world(x + w, y + h, u->screen_w, u->screen_h, u->z_plane, &x1, &y1);
    float z = u->z_plane;
    uint32_t base = (uint32_t)g_tvi;
    g_text_verts[g_tvi] = (UIVert){{x0,y0,z},{0,0,1},{u0,v0},{r,g,b,a}}; g_tvi++;
    g_text_verts[g_tvi] = (UIVert){{x1,y0,z},{0,0,1},{u1,v0},{r,g,b,a}}; g_tvi++;
    g_text_verts[g_tvi] = (UIVert){{x1,y1,z},{0,0,1},{u1,v1},{r,g,b,a}}; g_tvi++;
    g_text_verts[g_tvi] = (UIVert){{x0,y1,z},{0,0,1},{u0,v1},{r,g,b,a}}; g_tvi++;
    g_text_indices[g_tii++] = base; g_text_indices[g_tii++] = base+1; g_text_indices[g_tii++] = base+2;
    g_text_indices[g_tii++] = base; g_text_indices[g_tii++] = base+2; g_text_indices[g_tii++] = base+3;
    g_text_srv = srv;
}

static void flush_batch(ID3D11DeviceContext* ctx, UIVert* verts, uint32_t* indices,
                         int num_idx, ID3D11Buffer* vb, ID3D11Buffer* ib) {
    if (num_idx == 0) return;
    ctx->lpVtbl->UpdateSubresource(ctx, (ID3D11Resource*)vb, 0, NULL, verts, 0, 0);
    ctx->lpVtbl->UpdateSubresource(ctx, (ID3D11Resource*)ib, 0, NULL, indices, 0, 0);
    UINT stride = sizeof(UIVert), offset = 0;
    ctx->lpVtbl->IASetVertexBuffers(ctx, 0, 1, &vb, &stride, &offset);
    ctx->lpVtbl->IASetIndexBuffer(ctx, ib, DXGI_FORMAT_R32_UINT, 0);
    ctx->lpVtbl->DrawIndexed(ctx, (UINT)num_idx, 0, 0);
}

void uilayer_flush(UILayer* u, ID3D11DeviceContext* ctx) {
    if (g_ii == 0) return;
    Mat4 view = u->camera.view;
    Mat4 proj = u->camera.projection;
    Mat4 vp, vp_t;
    mat4_multiply(&view, &proj, &vp);
    mat4_transpose(&vp, &vp_t);

    D3D11_VIEWPORT vport = {0};
    vport.Width = u->screen_w; vport.Height = u->screen_h; vport.MaxDepth = 1.0f;
    ctx->lpVtbl->RSSetViewports(ctx, 1, &vport);
    ctx->lpVtbl->RSSetState(ctx, u->rstate);
    ctx->lpVtbl->OMSetDepthStencilState(ctx, u->ds_state, 0);
    ctx->lpVtbl->IASetPrimitiveTopology(ctx, D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
    shader_bind(u->shader, ctx);
    shader_set_matrix(u->shader, ctx, (float*)&vp_t);
    flush_batch(ctx, g_verts, g_indices, g_ii, u->vb, u->ib);
}

void uilayer_flush_text(UILayer* u, ID3D11DeviceContext* ctx) {
    if (g_tii == 0 || !g_text_srv || !u->text_vs || !u->text_ps) return;
    Mat4 view = u->camera.view;
    Mat4 proj = u->camera.projection;
    Mat4 vp, vp_t;
    mat4_multiply(&view, &proj, &vp);
    mat4_transpose(&vp, &vp_t);

    D3D11_VIEWPORT vport = {0};
    vport.Width = u->screen_w; vport.Height = u->screen_h; vport.MaxDepth = 1.0f;
    ctx->lpVtbl->RSSetViewports(ctx, 1, &vport);
    ctx->lpVtbl->RSSetState(ctx, u->rstate);
    ctx->lpVtbl->OMSetDepthStencilState(ctx, u->ds_state, 0);
    ctx->lpVtbl->IASetPrimitiveTopology(ctx, D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);

    ctx->lpVtbl->VSSetShader(ctx, u->text_vs, NULL, 0);
    ctx->lpVtbl->PSSetShader(ctx, u->text_ps, NULL, 0);
    ctx->lpVtbl->IASetInputLayout(ctx, u->text_layout);
    ctx->lpVtbl->UpdateSubresource(ctx, (ID3D11Resource*)u->shader->cb, 0, NULL, &vp_t, 0, 0);
    ctx->lpVtbl->VSSetConstantBuffers(ctx, 0, 1, &u->shader->cb);

    float bf[4] = {0};
    ctx->lpVtbl->OMSetBlendState(ctx, u->text_blend, bf, 0xFFFFFFFF);
    ctx->lpVtbl->PSSetShaderResources(ctx, 0, 1, &g_text_srv);
    ctx->lpVtbl->PSSetSamplers(ctx, 0, 1, &u->text_sampler);

    flush_batch(ctx, g_text_verts, g_text_indices, g_tii, u->text_vb, u->text_ib);
}