#include <Core/Common/String.h>
#include <Core/HAL/IO/IOStream.h>
#include <Core/Logging/LogManager.h>
#include <Core/Logging/LogTypes.h>
#include <Core/Memory/Allocator.h>
#include <Core/Memory/MemoryArena.h>
#include <ShaderCompiler.h>

internal void PrintHelp(void)
{
    LogMessage(LogCategory_Tool, "Usage: JulietShaderCompiler.exe <input> [options]");
    LogMessage(LogCategory_Tool, "Required Parameters:");
    const int kColumnWidth = 32;
    LogMessage(LogCategory_Tool, "  %-*s %s", kColumnWidth, "-o | --output <value>", "Output file.");
}

int main(int argc, char* argv[])
{
    int        exitCode       = 0;
    String     filename       = { 0 };
    ByteBuffer fileData       = { 0 };
    String     outputFilename = { 0 };
    String     entryPoint     = WrapString("main");

    bool8          sourceLangDetected      = false;
    SourceLanguage srcLanguage             = SourceLanguage_Invalid;
    bool8          dstShaderFormatDetected = false;
    ShaderFormat   dstShaderFormat         = ShaderFormat_Invalid;
    bool8          shaderStageDetected     = false;
    ShaderStage    shaderStage             = ShaderStage_Invalid;

    IOStream* outStream = NULL;
    Arena*    arena     = ArenaAllocate((ArenaParams){ .Name = "ShaderCompilerArena" });

    for (int idx = 1; idx < argc; ++idx)
    {
        String arg = WrapString(argv[idx]);
        if (CStr(arg)[0] == '-')
        {
            if (StringCompareCaseInsensitive(arg, ConstString("-h")) == 0 ||
                StringCompareCaseInsensitive(arg, ConstString("--help")) == 0)
            {
                PrintHelp();
                exitCode = 1;
                break;
            }

            if (StringCompareCaseInsensitive(arg, ConstString("-o")) == 0 ||
                StringCompareCaseInsensitive(arg, ConstString("--output")) == 0)
            {
                if (idx + 1 >= argc)
                {
                    LogError(LogCategory_Tool, "%s requires an argument", argv[idx]);
                    PrintHelp();
                    exitCode = 1;
                    break;
                }
                idx += 1;
                outputFilename = WrapString(argv[idx]);
            }
        }
        else if (!IsValidString(filename))
        {
            filename = arg;
        }
        else
        {
            LogError(LogCategory_Tool, "%s: Unknown parameter: %s", argv[0], CStr(arg));
            PrintHelp();
            exitCode = 1;
            break;
        }
    }

    if (exitCode == 0)
    {
        if (!IsValidString(filename))
        {
            LogError(LogCategory_Tool, "%s: missing input filename path", argv[0]);
            PrintHelp();
            exitCode = 1;
        }
        else if (!IsValidString(outputFilename))
        {
            LogError(LogCategory_Tool, "%s: missing output path", argv[0]);
            PrintHelp();
            exitCode = 1;
        }
    }

    if (exitCode == 0)
    {
        fileData = LoadFile(arena, filename);
        if (!IsValidByteBuffer(fileData))
        {
            LogError(LogCategory_Tool, "Invalid file: %s", CStr(filename));
            exitCode = 1;
        }
    }

    if (exitCode == 0)
    {
        if (!sourceLangDetected)
        {
            if (IsValidString(FindString(filename, WrapString("hlsl"))))
            {
                srcLanguage = SourceLanguage_HLSL;
            }
            StaticAssert(SourceLanguage_Count == 1, SourceLanguageCountCheck);
        }

        if (srcLanguage == SourceLanguage_Invalid)
        {
            LogError(LogCategory_Tool,
                     "%s: Invalid source language. Please either have a valid file extension or use -s/--src", argv[0]);
            PrintHelp();
            exitCode = 1;
        }
    }

    if (exitCode == 0)
    {
        if (!dstShaderFormatDetected)
        {
            if (IsValidString(FindString(outputFilename, WrapString("dxil"))))
            {
                dstShaderFormat = ShaderFormat_DXIL;
            }
            StaticAssert(ShaderFormat_Count == 1, ShaderFormatCountCheck);

            if (dstShaderFormat == ShaderFormat_Invalid)
            {
                LogError(LogCategory_Tool,
                         "%s: Invalid shader format. Please either have a valid file extension or use -d/--dest", argv[0]);
                PrintHelp();
                exitCode = 1;
            }
        }
    }

    if (exitCode == 0)
    {
        if (!shaderStageDetected)
        {
            if (IsValidString(FindString(filename, WrapString("vert"))))
            {
                shaderStage = ShaderStage_Vertex;
            }
            else if (IsValidString(FindString(filename, WrapString("frag"))))
            {
                shaderStage = ShaderStage_Fragment;
            }
            else if (IsValidString(FindString(filename, WrapString("comp"))))
            {
                shaderStage = ShaderStage_Compute;
            }
            StaticAssert(ShaderFormat_Count == 1, ShaderStageFormatCountCheck);

            if (shaderStage == ShaderStage_Invalid)
            {
                LogError(LogCategory_Tool, "%s: Invalid shader stage. Please include it in the filename", argv[0]);
                PrintHelp();
                exitCode = 1;
            }
        }
    }

    if (exitCode == 0)
    {
        outStream = IOFromFile(arena, outputFilename, WrapString("wb"));
        if (outStream == NULL)
        {
            LogError(LogCategory_Tool, "%s: Failed to open %s for writing", argv[0], CStr(outputFilename));
            exitCode = 1;
        }
    }

    if (exitCode == 0)
    {
        switch (srcLanguage)
        {
        case SourceLanguage_HLSL:
            {
                HLSLShaderInfo hlslInfo = { 0 };
                hlslInfo.ByteCodeBuffer = fileData;
                hlslInfo.EntryPoint     = entryPoint;
                hlslInfo.Name           = filename;
                hlslInfo.Stage          = shaderStage;

                switch (dstShaderFormat)
                {
                case ShaderFormat_DXIL:
                    {
                        ByteBuffer compiledBuffer = CompileDXIL(&hlslInfo);
                        if (!IsValidByteBuffer(compiledBuffer))
                        {
                            LogError(LogCategory_Tool, "Failed to compile DXBC from HLSL. Error: %s", "Unknown");
                        }
                        else
                        {
                            IOWrite(outStream, compiledBuffer);
                            FreeByteBuffer(&compiledBuffer);
                        }
                        break;
                    }
                default:
                    {
                        Assert(false && "Shader format is invalid");
                        exitCode = 1;
                        break;
                    }
                }
                break;
            }
        default:
            {
                Assert(false && "Source Language is invalid");
                exitCode = 1;
                break;
            }
        }
    }

    if (outStream != NULL)
    {
        IOClose(outStream);
    }
    ArenaRelease(arena);

    return exitCode;
}