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

static const char* default_vs =
    "struct VS_IN{float3 pos:POSITION;float3 n:NORMAL;float2 uv:TEXCOORD;float4 col:COLOR;};"
    "struct PS_IN{float4 pos:SV_POSITION;float3 world_n:NORMAL;float2 uv:TEXCOORD;float4 col:COLOR;};"
    "cbuffer CB:register(b0){float4x4 World;float4x4 WVP;float4 LightDir;float4 LightColor;float4 AmbientColor;float4 UseTexture;float4 DiffuseColor;};"
    "PS_IN main(VS_IN i){PS_IN o;o.pos=mul(float4(i.pos,1),WVP);o.world_n=normalize(mul(float4(i.n,0),World).xyz);o.uv=i.uv;o.col=i.col*DiffuseColor;return o;}";

static const char* default_ps =
    "struct PS_IN{float4 pos:SV_POSITION;float3 world_n:NORMAL;float2 uv:TEXCOORD;float4 col:COLOR;};"
    "cbuffer CB:register(b0){float4x4 World;float4x4 WVP;float4 LightDir;float4 LightColor;float4 AmbientColor;float4 UseTexture;float4 DiffuseColor;};"
    "Texture2D tex:register(t0);SamplerState samp:register(s0);"
    "float4 main(PS_IN i):SV_TARGET{float3 ld=normalize(LightDir.xyz);float diff=saturate(dot(i.world_n,ld));float4 light=AmbientColor+LightColor*diff;float4 albedo=i.col;if(UseTexture.x>0.5)albedo=albedo*tex.Sample(samp,i.uv);return albedo*light;}";

static ID3DBlob* compile_shader(const char* code, uint32_t len, const char* target) {
    ID3DBlob *blob = NULL, *err = NULL;
    UINT flags = D3DCOMPILE_OPTIMIZATION_LEVEL3 | D3DCOMPILE_PACK_MATRIX_ROW_MAJOR;
    D3DCompile(code, len, NULL, NULL, NULL, "main", target, flags, 0, &blob, &err);
    if (err) {
        const char* msg = (const char*)err->lpVtbl->GetBufferPointer(err);
        LOG_ERROR("Shader compile error (%s): %s", target, msg);
        err->lpVtbl->Release(err);
    }
    return blob;
}

Shader* shader_create(void) {
    return (Shader*)calloc(1, sizeof(Shader));
}

void shader_destroy(Shader* s) {
    if (!s) return;
    if (s->cb) s->cb->lpVtbl->Release(s->cb);
    if (s->layout) s->layout->lpVtbl->Release(s->layout);
    if (s->ps) s->ps->lpVtbl->Release(s->ps);
    if (s->vs) s->vs->lpVtbl->Release(s->vs);
    free(s);
}

int shader_compile(Shader* s, ID3D11Device* dev, const char* vs_code, const char* ps_code, uint32_t vs_len, uint32_t ps_len) {
    if (!vs_code || !ps_code) {
        vs_code = default_vs; ps_code = default_ps;
        vs_len = (uint32_t)strlen(default_vs); ps_len = (uint32_t)strlen(default_ps);
    }
    ID3DBlob* vs_blob = compile_shader(vs_code, vs_len, "vs_5_0");
    ID3DBlob* ps_blob = compile_shader(ps_code, ps_len, "ps_5_0");
    if (!vs_blob || !ps_blob) {
        LOG_ERROR("Failed to compile default shader (vs=%p ps=%p)", (void*)vs_blob, (void*)ps_blob);
        if (vs_blob) vs_blob->lpVtbl->Release(vs_blob);
        if (ps_blob) ps_blob->lpVtbl->Release(ps_blob);
        return 0;
    }
    dev->lpVtbl->CreateVertexShader(dev, vs_blob->lpVtbl->GetBufferPointer(vs_blob), vs_blob->lpVtbl->GetBufferSize(vs_blob), NULL, &s->vs);
    dev->lpVtbl->CreatePixelShader(dev, ps_blob->lpVtbl->GetBufferPointer(ps_blob), ps_blob->lpVtbl->GetBufferSize(ps_blob), NULL, &s->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, vs_blob->lpVtbl->GetBufferPointer(vs_blob), vs_blob->lpVtbl->GetBufferSize(vs_blob), &s->layout);
    D3D11_BUFFER_DESC bd = {0};
    bd.ByteWidth = 256; bd.Usage = D3D11_USAGE_DEFAULT; bd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    dev->lpVtbl->CreateBuffer(dev, &bd, NULL, &s->cb);
    vs_blob->lpVtbl->Release(vs_blob);
    ps_blob->lpVtbl->Release(ps_blob);
    return 1;
}

void shader_bind(Shader* s, ID3D11DeviceContext* ctx) {
    ctx->lpVtbl->VSSetShader(ctx, s->vs, NULL, 0);
    ctx->lpVtbl->PSSetShader(ctx, s->ps, NULL, 0);
    ctx->lpVtbl->IASetInputLayout(ctx, s->layout);
    ctx->lpVtbl->VSSetConstantBuffers(ctx, 0, 1, &s->cb);
    ctx->lpVtbl->PSSetConstantBuffers(ctx, 0, 1, &s->cb);
}

void shader_set_matrix(Shader* s, ID3D11DeviceContext* ctx, const float* data) {
    ctx->lpVtbl->UpdateSubresource(ctx, (ID3D11Resource*)s->cb, 0, NULL, data, 0, 0);
}