#include <d3dcompiler.h>

#include "shader_compiler.h"

extern "C" rt_shader_bytecode CompileDX11Shader(rt_shader_stage stage,
                                                rt_shader_optimization_level optimization,
                                                rt_text_span code,
                                                const char *file_path,
                                                rt_arena *arena) {
    rt_shader_bytecode bytecode = {};

    D3D_SHADER_MACRO defines[] = {
        {"RT_DX11",  "1"},
        {     NULL, NULL},
    };

    LPCSTR target = NULL;
    LPCSTR entrypoint = NULL;
    switch (stage) {
    case RT_SHADER_STAGE_VERTEX:
        target = "vs_5_0";
        entrypoint = "VsMain";
        break;
    case RT_SHADER_STAGE_FRAGMENT:
        target = "ps_5_0";
        entrypoint = "PsMain";
        break;
    case RT_SHADER_STAGE_COMPUTE:
        target = "cs_5_0";
        entrypoint = "CsMain";
        break;
    default:
        rtLog("AC", "Tried to compile an invalid shader stage %u for %s", stage, file_path);
        return bytecode;
    }

    UINT flags = D3DCOMPILE_PARTIAL_PRECISION | D3DCOMPILE_WARNINGS_ARE_ERRORS;
    switch (optimization) {
    case RT_SHADER_OPTIMIZATION_NONE:
        flags |= D3DCOMPILE_SKIP_OPTIMIZATION | D3DCOMPILE_OPTIMIZATION_LEVEL0;
        break;
    case RT_SHADER_OPTIMIZATION_SIZE:
        flags |= D3DCOMPILE_OPTIMIZATION_LEVEL2;
        break;
    case RT_SHADER_OPTIMIZATION_SPEED:
        flags |= D3DCOMPILE_OPTIMIZATION_LEVEL3;
        break;
    }
#ifdef RT_DEBUG
    flags |= D3DCOMPILE_DEBUG;
#endif
    ID3DBlob *bytes = nullptr, *error = nullptr;
    if (FAILED(D3DCompile2(code.start,
                           code.length,
                           file_path,
                           defines,
                           D3D_COMPILE_STANDARD_FILE_INCLUDE,
                           entrypoint,
                           target,
                           flags,
                           0,
                           0,
                           nullptr,
                           0,
                           &bytes,
                           &error))) {
        if (error) {
            rtLog("AC",
                  "Shader compilation failed (%s): %s",file_path,
                  (const char *)error->GetBufferPointer());
            error->Release();
            return bytecode;
        } else {
            rtLog("AC", "Shader compilation failed (%s): NO ERROR INFORMATION", file_path);
            return bytecode;
        }
    }

    bytecode.bytes = rtArenaPush(arena, bytes->GetBufferSize());
    if (bytecode.bytes) {
        bytecode.len = bytes->GetBufferSize();
        memcpy(bytecode.bytes, bytes->GetBufferPointer(), bytecode.len);
    } else {
        rtLog("AC", "Out of memory while compiling %s", file_path);
    }
    bytes->Release();

    return bytecode;
}