#include "graphics_device.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

GraphicsDevice* graphics_device_create(void) {
    return (GraphicsDevice*)calloc(1, sizeof(GraphicsDevice));
}

void graphics_device_destroy(GraphicsDevice* d) {
    if (!d) return;
    if (d->depth_stencil) d->depth_stencil->lpVtbl->Release(d->depth_stencil);
    if (d->depth_texture) d->depth_texture->lpVtbl->Release(d->depth_texture);
    if (d->render_target) d->render_target->lpVtbl->Release(d->render_target);
    if (d->swap_chain) d->swap_chain->lpVtbl->Release(d->swap_chain);
    if (d->context) d->context->lpVtbl->Release(d->context);
    if (d->device) d->device->lpVtbl->Release(d->device);
    free(d);
}

static int create_depth_stencil(GraphicsDevice* d, uint32_t w, uint32_t h) {
    D3D11_TEXTURE2D_DESC desc = {0};
    desc.Width = w; desc.Height = h;
    desc.MipLevels = 1; desc.ArraySize = 1;
    desc.Format = DXGI_FORMAT_D24_UNORM_S8_UINT;
    desc.SampleDesc.Count = 1;
    desc.Usage = D3D11_USAGE_DEFAULT;
    desc.BindFlags = D3D11_BIND_DEPTH_STENCIL;
    HRESULT hr = d->device->lpVtbl->CreateTexture2D(d->device, &desc, NULL, &d->depth_texture);
    if (FAILED(hr)) return 0;
    D3D11_DEPTH_STENCIL_VIEW_DESC dsv = {0};
    dsv.Format = desc.Format;
    dsv.ViewDimension = D3D11_DSV_DIMENSION_TEXTURE2D;
    hr = d->device->lpVtbl->CreateDepthStencilView(d->device, (ID3D11Resource*)d->depth_texture, &dsv, &d->depth_stencil);
    return SUCCEEDED(hr);
}

int graphics_device_init(GraphicsDevice* d, void* hwnd, uint32_t w, uint32_t h, int vsync) {
    d->width = w; d->height = h; d->vsync = vsync;
    DXGI_SWAP_CHAIN_DESC scd = {0};
    scd.BufferCount = 1;
    scd.BufferDesc.Width = w;
    scd.BufferDesc.Height = h;
    scd.BufferDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    scd.BufferDesc.RefreshRate.Numerator = 60;
    scd.BufferDesc.RefreshRate.Denominator = 1;
    scd.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
    scd.OutputWindow = (HWND)hwnd;
    scd.SampleDesc.Count = 1;
    scd.Windowed = TRUE;
    scd.SwapEffect = DXGI_SWAP_EFFECT_DISCARD;
    D3D_FEATURE_LEVEL levels[] = {D3D_FEATURE_LEVEL_11_0};
    UINT flags = 0;
#ifdef _DEBUG
    flags |= D3D11_CREATE_DEVICE_DEBUG;
#endif
    HRESULT hr = D3D11CreateDeviceAndSwapChain(NULL, D3D_DRIVER_TYPE_HARDWARE, NULL, flags, levels, 1, D3D11_SDK_VERSION, &scd, &d->swap_chain, &d->device, NULL, &d->context);
    if (FAILED(hr)) return 0;
    ID3D11Texture2D* back_buffer = NULL;
    d->swap_chain->lpVtbl->GetBuffer(d->swap_chain, 0, &IID_ID3D11Texture2D, (void**)&back_buffer);
    d->device->lpVtbl->CreateRenderTargetView(d->device, (ID3D11Resource*)back_buffer, NULL, &d->render_target);
    if (back_buffer) back_buffer->lpVtbl->Release(back_buffer);
    create_depth_stencil(d, w, h);
    d->context->lpVtbl->OMSetRenderTargets(d->context, 1, &d->render_target, d->depth_stencil);
    D3D11_VIEWPORT vp = {0};
    vp.Width = (float)w; vp.Height = (float)h;
    vp.MaxDepth = 1.0f;
    d->context->lpVtbl->RSSetViewports(d->context, 1, &vp);

    // 设置默认深度状态 (避免 UILayer init 破坏)
    D3D11_DEPTH_STENCIL_DESC dsd = {0};
    dsd.DepthEnable = TRUE;
    dsd.DepthWriteMask = D3D11_DEPTH_WRITE_MASK_ALL;
    dsd.DepthFunc = D3D11_COMPARISON_LESS_EQUAL;
    ID3D11DepthStencilState* default_dss = NULL;
    d->device->lpVtbl->CreateDepthStencilState(d->device, &dsd, &default_dss);
    if (default_dss) {
        d->context->lpVtbl->OMSetDepthStencilState(d->context, default_dss, 0);
        default_dss->lpVtbl->Release(default_dss);
    }

    // 设置默认光栅化状态：背面剔除 + 实体填充
    D3D11_RASTERIZER_DESC rd = {0};
    rd.FillMode = D3D11_FILL_SOLID;
    rd.CullMode = D3D11_CULL_BACK;
    rd.FrontCounterClockwise = FALSE;
    rd.DepthClipEnable = TRUE;
    ID3D11RasterizerState* default_rs = NULL;
    d->device->lpVtbl->CreateRasterizerState(d->device, &rd, &default_rs);
    if (default_rs) {
        d->context->lpVtbl->RSSetState(d->context, default_rs);
        default_rs->lpVtbl->Release(default_rs);
    }
    return 1;
}

void graphics_device_resize(GraphicsDevice* d, uint32_t w, uint32_t h) {
    if (w == 0 || h == 0) return;
    d->context->lpVtbl->OMSetRenderTargets(d->context, 0, NULL, NULL);
    if (d->render_target) { d->render_target->lpVtbl->Release(d->render_target); d->render_target = NULL; }
    if (d->depth_stencil) { d->depth_stencil->lpVtbl->Release(d->depth_stencil); d->depth_stencil = NULL; }
    if (d->depth_texture) { d->depth_texture->lpVtbl->Release(d->depth_texture); d->depth_texture = NULL; }
    d->swap_chain->lpVtbl->ResizeBuffers(d->swap_chain, 0, w, h, DXGI_FORMAT_UNKNOWN, 0);
    ID3D11Texture2D* back_buffer = NULL;
    d->swap_chain->lpVtbl->GetBuffer(d->swap_chain, 0, &IID_ID3D11Texture2D, (void**)&back_buffer);
    d->device->lpVtbl->CreateRenderTargetView(d->device, (ID3D11Resource*)back_buffer, NULL, &d->render_target);
    if (back_buffer) back_buffer->lpVtbl->Release(back_buffer);
    create_depth_stencil(d, w, h);
    d->context->lpVtbl->OMSetRenderTargets(d->context, 1, &d->render_target, d->depth_stencil);
    D3D11_VIEWPORT vp = {0};
    vp.Width = (float)w; vp.Height = (float)h; vp.MaxDepth = 1.0f;
    d->context->lpVtbl->RSSetViewports(d->context, 1, &vp);
    d->width = w; d->height = h;
}

void graphics_device_clear(GraphicsDevice* d, float r, float g, float b, float a) {
    float color[4] = {r, g, b, a};
    d->context->lpVtbl->ClearRenderTargetView(d->context, d->render_target, color);
    d->context->lpVtbl->ClearDepthStencilView(d->context, d->depth_stencil, D3D11_CLEAR_DEPTH | D3D11_CLEAR_STENCIL, 1.0f, 0);
}

void graphics_device_present(GraphicsDevice* d) {
    d->swap_chain->lpVtbl->Present(d->swap_chain, d->vsync ? 1 : 0, 0);
}

int graphics_device_save_screenshot(GraphicsDevice* d, const char* path) {
    if (!d || !d->swap_chain || !d->device || !d->context) return 0;

    ID3D11Texture2D* back = NULL;
    HRESULT hr = d->swap_chain->lpVtbl->GetBuffer(d->swap_chain, 0, &IID_ID3D11Texture2D, (void**)&back);
    if (FAILED(hr) || !back) return 0;

    D3D11_TEXTURE2D_DESC desc = {0};
    back->lpVtbl->GetDesc(back, &desc);

    D3D11_TEXTURE2D_DESC sd = desc;
    sd.BindFlags = 0;
    sd.Usage = D3D11_USAGE_STAGING;
    sd.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
    sd.MiscFlags = 0;

    ID3D11Texture2D* staging = NULL;
    hr = d->device->lpVtbl->CreateTexture2D(d->device, &sd, NULL, &staging);
    if (FAILED(hr) || !staging) { back->lpVtbl->Release(back); return 0; }

    d->context->lpVtbl->CopyResource(d->context, (ID3D11Resource*)staging, (ID3D11Resource*)back);

    D3D11_MAPPED_SUBRESOURCE map = {0};
    hr = d->context->lpVtbl->Map(d->context, (ID3D11Resource*)staging, 0, D3D11_MAP_READ, 0, &map);
    if (FAILED(hr)) { staging->lpVtbl->Release(staging); back->lpVtbl->Release(back); return 0; }

    FILE* f = fopen(path, "wb");
    if (!f) { d->context->lpVtbl->Unmap(d->context, (ID3D11Resource*)staging, 0); staging->lpVtbl->Release(staging); back->lpVtbl->Release(back); return 0; }

    int w = (int)desc.Width;
    int h = (int)desc.Height;
    int row_size = ((w * 3 + 3) / 4) * 4;
    int file_size = 54 + row_size * h;

    unsigned char header[54] = {0};
    header[0] = 'B'; header[1] = 'M';
    memcpy(header + 2, &file_size, 4);
    header[10] = 54;
    header[14] = 40;
    memcpy(header + 18, &w, 4);
    memcpy(header + 22, &h, 4);
    header[26] = 1;
    header[28] = 24;
    fwrite(header, 1, 54, f);

    unsigned char* row = (unsigned char*)malloc(row_size);
    for (int y = 0; y < h; y++) {
        uint8_t* src = (uint8_t*)map.pData + (h - 1 - y) * map.RowPitch;
        for (int x = 0; x < w; x++) {
            // 后台缓冲区为 RGBA，BMP 24-bit 需要 BGR
            row[x * 3 + 0] = src[x * 4 + 2];
            row[x * 3 + 1] = src[x * 4 + 1];
            row[x * 3 + 2] = src[x * 4 + 0];
        }
        fwrite(row, 1, row_size, f);
    }
    free(row);
    fclose(f);

    d->context->lpVtbl->Unmap(d->context, (ID3D11Resource*)staging, 0);
    staging->lpVtbl->Release(staging);
    back->lpVtbl->Release(back);
    return 1;
}