#include <ShaderCompiler.h>

#include <Core/Common/CoreUtils.h>
#include <Core/Logging/LogManager.h>
#include <Core/Logging/LogTypes.h>
#include <Core/Memory/Allocator.h>
#include <DXShaderCompiler/DXCompiler.h>

internal ByteBuffer CompileWithDXC(const HLSLShaderInfo* info, bool8 compileToSpirv)
{
    Assert(info != NULL);
    Assert(!compileToSpirv && "SPIRV is not supported yet");

    ByteBuffer     outBuffer = { 0 };
    IDxcCompiler3* compiler  = NULL;
    DxcCreateInstance(&CLSID_DxcCompiler, IID_IDxcCompiler3, (void**)&compiler);

    IDxcUtils* utils = NULL;
    DxcCreateInstance(&CLSID_DxcUtils, IID_IDxcUtils, (void**)&utils);

    IDxcResult* dxcResult = NULL;
    LPCWSTR*    args      = NULL;

    if (compiler == NULL)
    {
        LogError(LogCategory_Tool, "Cannot create DXCompiler instance");
    }
    else if (utils == NULL)
    {
        LogError(LogCategory_Graphics, "Cannot create IDxcUtils instance");
    }
    else
    {
        IDxcIncludeHandler* includeHandler = NULL;
        utils->lpVtbl->CreateDefaultIncludeHandler(utils, &includeHandler);
        if (includeHandler == NULL)
        {
            LogError(LogCategory_Graphics, "Cannot create Default include handler");
        }
        else
        {
            StringBuffer convertedBuffer    = { 0 };
            char         inplaceBuffer[1024];
            convertedBuffer.Str             = inplaceBuffer;
            convertedBuffer.Capacity        = sizeof(inplaceBuffer);
            const bool8 kAddNullTerminator  = true;
            if (!ConvertString(StringEncoding_UTF8, StringEncoding_UTF16, info->EntryPoint, &convertedBuffer, kAddNullTerminator))
            {
                LogError(LogCategory_Tool, "Cannot convert entry point string to utf16");
            }
            else
            {
                const size_t kExpectedArgCount = 32;
                uint32       argCount          = 0;
                args                           = (LPCWSTR*)Calloc(kExpectedArgCount, sizeof(LPCWSTR));
                if (args == NULL)
                {
                    LogError(LogCategory_Tool, "Couldn't allocate Args array");
                }
                else
                {
                    args[argCount++] = L"-E";
                    args[argCount++] = (LPCWSTR)inplaceBuffer;

                    DxcBuffer source;
                    source.Ptr      = info->ByteCodeBuffer.Data;
                    source.Size     = info->ByteCodeBuffer.Size;
                    source.Encoding = DXC_CP_ACP;

                    bool8 stageValid = true;
                    switch (info->Stage)
                    {
                    case ShaderStage_Vertex:
                        {
                            args[argCount++] = L"-T";
                            args[argCount++] = L"vs_6_6";
                            break;
                        }
                    case ShaderStage_Fragment:
                        {
                            args[argCount++] = L"-T";
                            args[argCount++] = L"ps_6_6";
                            break;
                        }
                    case ShaderStage_Compute:
                        {
                            args[argCount++] = L"-T";
                            args[argCount++] = L"cs_6_6";
                            break;
                        }
                    default:
                        {
                            Assert(false && "Invalid shader stage");
                            stageValid = false;
                            break;
                        }
                    }

                    if (stageValid)
                    {
                        if (compileToSpirv)
                        {
                            args[argCount++] = L"-spirv";
                            args[argCount++] = L"-fspv-flatten-resource-arrays";
                        }

                        char nameInplaceBuffer[1024];
                        if (IsValidString(info->Name))
                        {
                            StringBuffer nameUtf16 = { 0 };
                            nameUtf16.Str          = nameInplaceBuffer;
                            nameUtf16.Capacity     = sizeof(nameInplaceBuffer);
                            if (ConvertString(StringEncoding_UTF8, StringEncoding_UTF16, info->Name, &nameUtf16, kAddNullTerminator))
                            {
                                args[argCount++] = (LPCWSTR)nameInplaceBuffer;
                            }
                        }

                        HRESULT result = compiler->lpVtbl->Compile(compiler, &source, args, argCount, includeHandler,
                                                                   IID_IDxcResult, (void**)&dxcResult);
                        if (result < 0)
                        {
                            LogError(LogCategory_Tool, "IDxcShaderCompiler3::Compile failed: %X", result);
                        }
                        else if (dxcResult == NULL)
                        {
                            LogError(LogCategory_Tool, "%s", "HLSL compilation failed with no IDxcResult");
                        }
                        else
                        {
                            IDxcBlob*     blob   = NULL;
                            IDxcBlobUtf8* errors = NULL;
                            result = dxcResult->lpVtbl->GetOutput(dxcResult, DXC_OUT_OBJECT, IID_IDxcBlob, (void**)&blob, NULL);
                            if (result < 0)
                            {
                                dxcResult->lpVtbl->GetOutput(dxcResult, DXC_OUT_ERRORS, IID_IDxcBlobUtf8, (void**)&errors, NULL);
                                if (errors != NULL && errors->lpVtbl->GetBufferSize(errors) != 0)
                                {
                                    LogError(LogCategory_Tool, "HLSL compilation failed: %s",
                                             (char*)errors->lpVtbl->GetBufferPointer(errors));
                                }
                                else
                                {
                                    LogError(LogCategory_Tool, "Compilation failed with unknown error");
                                }
                            }
                            else
                            {
                                dxcResult->lpVtbl->GetOutput(dxcResult, DXC_OUT_ERRORS, IID_IDxcBlobUtf8, (void**)&errors, NULL);
                                if (errors != NULL && errors->lpVtbl->GetBufferSize(errors) != 0)
                                {
                                    LogWarning(LogCategory_Tool, "HLSL compiled with warnings: %s",
                                               (char*)errors->lpVtbl->GetBufferPointer(errors));
                                }

                                outBuffer.Size = blob->lpVtbl->GetBufferSize(blob);
                                outBuffer.Data = (Byte*)Malloc(outBuffer.Size);
                                MemCopy(outBuffer.Data, blob->lpVtbl->GetBufferPointer(blob), outBuffer.Size);

                                blob->lpVtbl->Release(blob);
                            }
                        }
                    }
                }
            }
        }
    }

    SafeFree(args);
    if (dxcResult != NULL)
    {
        dxcResult->lpVtbl->Release(dxcResult);
    }
    if (compiler != NULL)
    {
        compiler->lpVtbl->Release(compiler);
    }
    if (utils != NULL)
    {
        utils->lpVtbl->Release(utils);
    }

    return outBuffer;
}

ByteBuffer CompileDXIL(const HLSLShaderInfo* info)
{
    return CompileWithDXC(info, false);
}