#ifdef _WIN32
// Required by dxcapi.h, does not work with WIN32_LEAN_AND_MEAN
#include <Windows.h>
#endif
#if defined(__GNUC__) || defined(__clang__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wconversion"
#endif
#include <dxcapi.h>
#if defined(__GNUC__) || defined(__clang__)
#pragma GCC diagnostic pop
#endif
#include <stdlib.h>

#include "processor.h"
#include "shader_compiler.h"

#ifdef __linux__
#define WCHAR wchar_t
#endif

extern "C" rt_shader_bytecode CompileVulkanShader(rt_shader_stage stage,
                                                  rt_shader_optimization_level optimization,
                                                  rt_text_span code,
                                                  const char *file_path,
                                                  rt_arena *arena) {
    rt_shader_bytecode bc;
    memset(&bc, 0, sizeof(bc));

    // NOTE(Kevin): This looks like some sort of autodetect?
    // Maybe fix this to UTF8?
    uint32_t codepage = DXC_CP_ACP;

    // NOTE(Kevin): 6_1 is used at https://docs.vulkan.org/guide/latest/hlsl.html
    // Check if this is what we want.
    // For example: 6_2 is what allows the usage of 16 bit types
    LPCWSTR target_profile = nullptr;
    LPCWSTR entry          = nullptr;
    switch (stage) {
    case RT_SHADER_STAGE_VERTEX:
        target_profile = L"vs_6_1";
        entry          = L"VsMain";
        break;
    case RT_SHADER_STAGE_FRAGMENT:
        target_profile = L"ps_6_1";
        entry          = L"PsMain";
        break;
    case RT_SHADER_STAGE_COMPUTE:
        target_profile = L"cs_6_1";
        entry          = L"CsMain";
        break;
    default:
        rtReportError("AC", "Invalid shader stage %u for shader %s", stage, file_path);
        return bc;
    }

    LPCWSTR optimization_arg = nullptr;
    switch (optimization) {
    case RT_SHADER_OPTIMIZATION_NONE:
        optimization_arg = L"-Od";
        break;
    case RT_SHADER_OPTIMIZATION_SIZE:
    case RT_SHADER_OPTIMIZATION_SPEED:
        optimization_arg = L"-O3";
        break;
    }

    // Init DXC library and compiler
    HRESULT hr;
    IDxcLibrary *library;
    hr = DxcCreateInstance(CLSID_DxcLibrary, IID_PPV_ARGS(&library));
    if (FAILED(hr)) {
        rtReportError("AC", "Failed to init the DXC library.");
        return bc;
    }
    IDxcCompiler3 *compiler;
    hr = DxcCreateInstance(CLSID_DxcCompiler, IID_PPV_ARGS(&compiler));
    if (FAILED(hr)) {
        library->Release();
        rtReportError("AC", "Failed to init the DXC compiler.");
        return bc;
    }
    IDxcUtils *utils;
    hr = DxcCreateInstance(CLSID_DxcUtils, IID_PPV_ARGS(&utils));
    if (FAILED(hr)) {
        compiler->Release();
        library->Release();
        rtReportError("AC", "Failed to init the DXC utils.");
        return bc;
    }
    IDxcIncludeHandler *include_handler;
    hr = utils->CreateDefaultIncludeHandler(&include_handler);
    if (FAILED(hr)) {
        utils->Release();
        compiler->Release();
        library->Release();
        rtReportError("AC", "Failed to init the DXC include handler.");
        return bc;
    }

    // Gather arguments.
    // TODO(kevin): Maybe add -Qstrip_debug & -Qstrip_reflect?
    WCHAR wname[RT_PATH_MAX];
    rtUTF8ToWStr(file_path, wname, RT_PATH_MAX);
    // clang-format off
    LPCWSTR args[] = {
        wname, 
        L"-E", entry,
        L"-T", target_profile,
        optimization_arg,
        L"-D", L"RT_VULKAN",
        L"-spirv"
    };
    // clang-format on

    // Compile!
    DxcBuffer buffer{};
    buffer.Encoding = codepage;
    buffer.Ptr      = (LPVOID)code.start;
    buffer.Size     = (SIZE_T)code.length;
    IDxcResult *result;
    hr = compiler->Compile(&buffer,
                           args,
                           RT_ARRAY_COUNT(args),
                           include_handler,
                           IID_PPV_ARGS(&result));
    if (SUCCEEDED(hr)) {
        result->GetStatus(&hr);
    }
    if (FAILED(hr) && result) {
        // Error occured
        IDxcBlobEncoding *error_blob;
        hr = result->GetErrorBuffer(&error_blob);
        if (SUCCEEDED(hr) && error_blob && error_blob->GetBufferSize() > 0) {
            const char *msg = (const char *)error_blob->GetBufferPointer();
            rtLog("AC", "Shader %s compilation failed: %s", file_path, msg);
            error_blob->Release();
        } else {
            rtLog("AC", "Shader %s compilation failed. No error information available!");
        }
        include_handler->Release();
        utils->Release();
        compiler->Release();
        library->Release();
        return bc;
    }

    IDxcBlob *code_blob;
    result->GetResult(&code_blob);

    bc.bytes = rtArenaPush(arena, code_blob->GetBufferSize());
    if (bc.bytes) {
        bc.len = code_blob->GetBufferSize();
        memcpy(bc.bytes, code_blob->GetBufferPointer(), bc.len);
    } else {
        rtLog("AC", "Out of memory while compiling %s", file_path);
    }
    include_handler->Release();
    utils->Release();
    compiler->Release();
    library->Release();
    return bc;
}
