#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>

#pragma comment(lib, "dxcompiler.lib")

using namespace Juliet;

namespace
{
    ByteBuffer CompileWithDXC(HLSLShaderInfo& info, bool compileToSpirv)
    {
        Assert(!compileToSpirv && "SPIRV is not supported yet");

        IDxcCompiler3* compiler = nullptr;
        DxcCreateInstance(&CLSID_DxcCompiler, IID_IDxcCompiler3, reinterpret_cast<void**>(&compiler));

        IDxcUtils* utils = nullptr;
        DxcCreateInstance(&CLSID_DxcUtils, &IID_IDxcUtils, reinterpret_cast<void**>(&utils));

        if (compiler == nullptr)
        {
            LogError(LogCategory::Tool, "Cannot create DXCompiler instance");
            return {};
        }

        IDxcResult* dxcResult = nullptr;
        LPCWSTR*    args      = nullptr;
        auto        deferred  = Defer(
            [&]()
            {
                SafeFree(args);
                if (dxcResult)
                {
                    dxcResult->lpVtbl->Release(dxcResult);
                }
                if (compiler)
                {
                    compiler->lpVtbl->Release(compiler);
                }
                if (utils)
                {
                    utils->lpVtbl->Release(utils);
                }
            });

        if (utils == nullptr)
        {
            LogError(LogCategory::Graphics, "Cannot create IDxcUtils instance");
            return {};
        }

        IDxcIncludeHandler* includeHandler;
        utils->lpVtbl->CreateDefaultIncludeHandler(utils, &includeHandler);
        if (includeHandler == nullptr)
        {
            LogError(LogCategory::Graphics, "Cannot create Default include handler");
            return {};
        }

        StringBuffer convertedBuffer = {};
        char         inplaceBuffer[1024];
        convertedBuffer.Data     = inplaceBuffer;
        convertedBuffer.Capacity = sizeof(inplaceBuffer);
        if (!ConvertString(StringEncoding::UTF8, StringEncoding::UTF16, info.EntryPoint, convertedBuffer))
        {
            LogError(LogCategory::Tool, "Cannot convert entry point string to utf16");
            return {};
        }

        // Creating args for the compiler
        // Allocating enough to fit them all
        static constexpr size_t kExpectedArgCount = 32;
        uint32                  argCount          = 0;
        args                                      = static_cast<LPCWSTR*>(Calloc(kExpectedArgCount, sizeof(LPCWSTR)));
        if (!args)
        {
            LogError(LogCategory::Tool, "Couldn't allocate Args array");
            return {};
        }

        args[argCount++] = const_cast<LPCWSTR>(L"-E");
        args[argCount++] = reinterpret_cast<LPCWSTR>(inplaceBuffer);

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

        switch (info.Stage)
        {
        case ShaderStage::Vertex:
            {
                args[argCount++] = const_cast<LPCWSTR>(L"-T");
                args[argCount++] = const_cast<LPCWSTR>(L"vs_6_0");
                break;
            }
        case ShaderStage::Fragment:
            {
                args[argCount++] = const_cast<LPCWSTR>(L"-T");
                args[argCount++] = const_cast<LPCWSTR>(L"ps_6_0");
                break;
            }
        case ShaderStage::Compute:
            {
                args[argCount++] = const_cast<LPCWSTR>(L"-T");
                args[argCount++] = const_cast<LPCWSTR>(L"cs_6_0");
                break;
            }
        case ShaderStage::Invalid:
            {
                Assert(false && "Invalid shader stage");
                return {};
            }
        }

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

        char nameInplaceBuffer[1024];
        if (IsValid(info.Name))
        {
            StringBuffer nameUtf16 = {};
            nameUtf16.Data         = nameInplaceBuffer;
            nameUtf16.Capacity     = sizeof(nameInplaceBuffer);
            if (ConvertString(StringEncoding::UTF8, StringEncoding::UTF16, info.Name, nameUtf16))
            {
                args[argCount++] = reinterpret_cast<LPCWSTR>(nameInplaceBuffer);
            }
        }

        HRESULT result = compiler->lpVtbl->Compile(compiler, &source, args, argCount, includeHandler, IID_IDxcResult,
                                                   reinterpret_cast<void**>(&dxcResult));
        if (result < 0)
        {
            LogError(LogCategory::Tool, "IDxcShaderCompiler3::Compile failed: %X", result);
            return {};
        }

        if (dxcResult == nullptr)
        {
            LogError(LogCategory::Tool, "%s", "HLSL compilation failed with no IDxcResult");
            return {};
        }

        IDxcBlob*     blob   = nullptr;
        IDxcBlobUtf8* errors = nullptr;
        result = dxcResult->lpVtbl->GetOutput(dxcResult, DXC_OUT_OBJECT, IID_IDxcBlob, reinterpret_cast<void**>(&blob), nullptr);
        if (result < 0)
        {
            // Compilation failed, display errors
            dxcResult->lpVtbl->GetOutput(dxcResult, DXC_OUT_ERRORS, IID_IDxcBlobUtf8, reinterpret_cast<void**>(&errors), nullptr);
            if (errors != nullptr && errors->lpVtbl->GetBufferSize(errors) != 0)
            {
                LogError(LogCategory::Tool, "HLSL compilation failed: %s",
                         static_cast<char*>(errors->lpVtbl->GetBufferPointer(errors)));
            }
            else
            {
                LogError(LogCategory::Tool, "Compilation failed with unknown error");
            }

            return {};
        }

        // If compilation succeeded, but there are errors, those are warnings
        dxcResult->lpVtbl->GetOutput(dxcResult, DXC_OUT_ERRORS, IID_IDxcBlobUtf8, reinterpret_cast<void**>(&errors), nullptr);
        if (errors != nullptr && errors->lpVtbl->GetBufferSize(errors) != 0)
        {
            LogWarning(LogCategory::Tool, "HLSL compiled with warnings: %s",
                       static_cast<char*>(errors->lpVtbl->GetBufferPointer(errors)));
        }

        // TODO: Scratch allocator. No malloc
        ByteBuffer buffer = {};
        buffer.Size       = blob->lpVtbl->GetBufferSize(blob);
        buffer.Data       = static_cast<Byte*>(Malloc(buffer.Size));
        MemCopy(buffer.Data, blob->lpVtbl->GetBufferPointer(blob), buffer.Size);

        blob->lpVtbl->Release(blob);

        return buffer;
    }
} // namespace

ByteBuffer CompileDXIL(HLSLShaderInfo& info)
{
    // TODO: Use SPIRV an do hlsl -> spirv -> hlsl -> dxil.
    // This is from Shadercross SDL_ShaderCross_CompileDXILFromHLSL.
    // Not sure why yet.

    return CompileWithDXC(info, false);
}