#include <Core/Common/EnumUtils.h>
#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 <ShaderCompiler.h>

using namespace Juliet;

void PrintHelp()
{
    LogMessage(LogCategory::Tool, "Usage: JulietShaderCompiler.exe <input> [options]");
    LogMessage(LogCategory::Tool, "Required Parameters:");
    static constexpr int kColumnWdith = 32;
    LogMessage(LogCategory::Tool, "  %-*s %s", kColumnWdith, "-o | --output <value>", "Output file.");
    // LogMessage(LogCategory::Tool, "\n");
    // LogMessage(LogCategory::Tool, "Required Parameters that can be inferred from filename:");
    // LogMessage(LogCategory::Tool, "  %-*s %s", kColumnWdith, "-s | --src <value>", "Source Language. Values: [HLSL*] | * Default");
    // LogMessage(LogCategory::Tool, "  %-*s %s", kColumnWdith, "-d | --dest <value>", "Destination Format. Values: [DXIL*] | * Default");
}

int main(int argc, char* argv[])
{
    String     filename = {};
    ByteBuffer fileData = {};

    String outputFilename = {};

    String entryPoint = WrapString("main");

    bool sourceLangDetected = false;
    auto srcLanguage        = SourceLanguage::Invalid;

    bool dstShaderFormatDetected = false;
    auto dstShaderFormat         = ShaderFormat::Invalid;

    bool shaderStageDetected = false;
    auto shaderStage         = ShaderStage::Invalid;

    IOStream* outStream = nullptr;

    auto deferred = Defer(
        [&]()
        {
            if (IsValid(fileData))
            {
                Free(fileData);
            }
            if (outStream)
            {
                IOClose(outStream);
            }
            // Pause here to not close the console window immediately on stop while debugging
            __asm int 3;
        });

    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();
                return 1;
            }

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

    if (!IsValid(filename))
    {
        LogError(LogCategory::Tool, "%s: missing input filename path", argv[0]);
        PrintHelp();
        return 1;
    }
    if (!IsValid(outputFilename))
    {
        LogError(LogCategory::Tool, "%s: missing output path", argv[0]);
        PrintHelp();
        return 1;
    }

    fileData = LoadFile(filename);
    if (!IsValid(fileData))
    {
        LogError(LogCategory::Tool, "Invalid file: %s", CStr(filename));
        return 1;
    }

    // Find the source language from filename
    {
        if (!sourceLangDetected)
        {
            if (IsValid(FindString(filename, WrapString("hlsl"))))
            {
                srcLanguage = SourceLanguage::HLSL;
            }
            static_assert(ToUnderlying(SourceLanguage::Count) == 1, "SourceLanguage must be one of: HLSL");
        }

        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();
            return 1;
        }
    }

    // Find the destination format from output file name
    {
        if (!dstShaderFormatDetected)
        {
            if (IsValid(FindString(outputFilename, WrapString("dxil"))))
            {
                dstShaderFormat = ShaderFormat::DXIL;
            }
            static_assert(ToUnderlying(ShaderFormat::Count) == 1, "ShaderFormat must be one of: DXIL");

            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();
                return 1;
            }
        }
    }

    // Find shader stage from filename
    {
        if (!shaderStageDetected)
        {
            // TODO Need a case insensitive FindString
            if (IsValid(FindString(filename, WrapString("vert"))))
            {
                shaderStage = ShaderStage::Vertex;
            }
            else if (IsValid(FindString(filename, WrapString("frag"))))
            {
                shaderStage = ShaderStage::Fragment;
            }
            else if (IsValid(FindString(filename, WrapString("comp"))))
            {
                shaderStage = ShaderStage::Compute;
            }
            static_assert(ToUnderlying(ShaderFormat::Count) == 1, "ShaderStage must be one of: vert, frag, comp");

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

    outStream = IOFromFile(outputFilename, WrapString("w"));
    if (!outStream)
    {
        LogError(LogCategory::Tool, "%s: Failed to open %s for writing", argv[0], CStr(outputFilename));
        return 1;
    }

    switch (srcLanguage)
    {
    case SourceLanguage::HLSL:
        {
            HLSLShaderInfo hlslInfo = {};
            hlslInfo.ByteCodeBuffer = fileData;
            hlslInfo.EntryPoint     = entryPoint;
            hlslInfo.Name           = filename;
            hlslInfo.Stage          = shaderStage;

            switch (dstShaderFormat)
            {
            case ShaderFormat::DXIL:
                {
                    ByteBuffer compiledBuffer = CompileDXIL(hlslInfo);
                    if (!IsValid(compiledBuffer))
                    {
                        LogError(LogCategory::Tool, "Failed to compile DXBC from HLSL. Error: %s", "Unknown");
                    }
                    else
                    {
                        IOWrite(outStream, compiledBuffer);
                        Free(compiledBuffer);
                    }
                    break;
                }
            case ShaderFormat::Invalid: Assert(false && "Shader format is invalid"); return 1;
            }

            break;
        }
    case SourceLanguage::Invalid: Assert(false && "Source Language is invalid"); return 1;
    }

    return 0;
}