#include "shader.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;float4 col:COLOR;};"
    "cbuffer CB:register(b0){float4x4 WVP;};"
    "PS_IN main(VS_IN i){PS_IN o;o.pos=mul(float4(i.pos,1),WVP);o.col=i.col;return o;}";

static const char* default_ps =
    "struct PS_IN{float4 pos:SV_POSITION;float4 col:COLOR;};"
    "float4 main(PS_IN i):SV_TARGET{return i.col;}";

static ID3DBlob* compile_shader(const char* code, uint32_t len, const char* target) {
    ID3DBlob *blob = NULL, *err = NULL;
    D3DCompile(code, len, NULL, NULL, NULL, "main", target, D3DCOMPILE_OPTIMIZATION_LEVEL3, 0, &blob, &err);
    if (err) 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) {
        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);
}

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