#include "CppGen.h"
#include "Logging.h"
#include "TypeDef.h"
#include "bx/string.h"

#include <bx/file.h>
#include <cstdarg>
#include <memoryapi.h>
#include <winnt.h>

#define WIN32_LEAN_AND_MEAN
#include "Windows.h" // IWYU pragma: keep

namespace WriteTemplates
{
    constexpr char FileHeaderStart[] =
        R"END(#pragma once
#include <cstdint>

namespace Gen
{
    struct Serializer;
    struct Deserializer;
)END";
    constexpr char FileCppStart[] = R"END(#include "Def.h"
#include "Generated.h"

namespace Gen
{
)END";

    constexpr char StructHeader2[] =
        R"END(    struct %s
    {
        static constexpr uint16_t TypeIdx = %u;
)END";

    constexpr char StructField4[] =
        R"END(        %s %s%s = %s;
)END";

    constexpr char StructEnd[] =
        R"END(    };
)END";

    constexpr char EnumHeader4[] =
        R"END(    struct %s
    {
        static constexpr uint16_t TypeIdx = %u;
        static constexpr int32_t EntryCount = %u;
        enum Enum : %s
        {
)END";

    constexpr char EnumField1[] =
        R"END(            %s,
)END";

    constexpr char EnumNamesStart2[] =
        R"END(        };
        static constexpr char %s[EntryCount][%u]
        {
)END";

    constexpr char EnumNamesEntry1[] =
        R"END(            "%s",
)END";

    constexpr char EnumNamesEnd[] =
        R"END(        };
    };
)END";

    constexpr char FileEnd[] =
        R"END(}
)END";

    constexpr char SaveFuncHeader1[] =
        R"END(    bool Save(const %s* obj, uint32_t count, Serializer& serializer);
)END";
    constexpr char SaveFuncBodyStart1[] = R"END(    bool Save(const %s* obj, uint32_t count, Serializer& serializer)
    {
        bool isOk = true;
        for (uint32_t i = 0; i < count; ++i)
        {
)END";
    constexpr char SaveFuncBodyType3[] = R"END(            isOk = Save(%sobj[i].%s, %u, serializer) && isOk;
)END";
    constexpr char SaveFuncBodyNative[] = R"END(            isOk = serializer.Write(&obj[i], sizeof(obj[i])) && isOk;
)END";
    constexpr char SaveFuncBodyEnum1[] = R"END(            auto val = (%s)obj[i];
            isOk = Save(&val, 1, serializer) && isOk;
)END";
    constexpr char SaveFuncBodyEnd[] = R"END(        }
        return isOk;
    }
)END";

    constexpr char LoadFuncHeader1[] =
        R"END(    bool Load(%s* obj, uint32_t count, Deserializer& serializer);
)END";
    constexpr char LoadFuncBodyStart1[] = R"END(    bool Load(%s* obj, uint32_t count, Deserializer& serializer)
    {
)END";
    constexpr char LoadFuncBodySetupStart2[] =
        R"END(        const char* typeName = Meta::Metadata.TypeDefinitions[%s::TypeIdx].Name;

        // Quick match
        int32_t matchedHashIdx =
            serializer.TypeBuf.FindHash(Meta::Metadata.TypeDefinitions[%s::TypeIdx].Hash);
        if (matchedHashIdx >= 0)
        {
            assert(bx::strCmp(serializer.TypeBuf.Defs[matchedHashIdx].Name, typeName) == 0);
            bool isOk = true;
            for (uint32_t i = 0; i < count; ++i)
            {
)END";
    constexpr char LoadFuncBodyQuickLoad2[] = R"END(                isOk = Load(%sobj[i].%s, %u, serializer) && isOk;
)END";
    constexpr char LoadFuncBodySetupEnd[] = R"END(            }
            // if we're not ok here, something went really wrong
            assert(isOk);
            return isOk;
        }

        // Failed to resolve hash, the type definition chaned since the file was saved! try to match by name.
        int32_t nameMatchIdx = serializer.TypeBuf.FindDefByName(typeName);
        if (nameMatchIdx < 0)
        {
            // Name match failed, caller has to handle this and potentially skip some bytes
            return false;
        }

        // Successfully matched name, but we need to follow the definition of the file now!
        const Meta::TypeDef& matchedDef = serializer.TypeBuf.Defs[nameMatchIdx];

        // Figure out new member mapping
        uint64_t WriteDestinations[64];
        for (int32_t i = 0; i < BX_COUNTOF(WriteDestinations); ++i)
        {
            WriteDestinations[i] = UINT64_MAX;
        }
        for (uint32_t i = 0; i < matchedDef.ChildCount; ++i)
        {
            const bx::StringView memberName = {&serializer.MemberNameBuf[matchedDef.MemberNameIndices[i].Offset], matchedDef.MemberNameIndices[i].Size};
            const char* memberTypeName = serializer.TypeBuf.Defs[matchedDef.ChildIndices[i]].Name;
)END";
    constexpr char LoadFuncBodyMemberCheck4[] =
        R"END(            if (bx::strCmp(memberName, "%s") == 0 && bx::strCmp(memberTypeName, "%s") == 0)
            {
                WriteDestinations[i] = offsetof(%s, %s);
            }
)END";
    constexpr char LoadFuncBodyTypeUpgradeStart[] = R"END(        }

        // Start reading in file order, skipping things that we don't know by name and type
        bool isOk = true;
        for (uint32_t i = 0; i < count; ++i)
        {
            uint8_t* objBasePtr = reinterpret_cast<uint8_t*>(&obj[i]);

            for (uint32_t j = 0; j < matchedDef.ChildCount; ++j)
            {
                const Meta::TypeDef& childDef = serializer.TypeBuf.Defs[matchedDef.ChildIndices[j]];
                const bx::StringView memberName = {&serializer.MemberNameBuf[matchedDef.MemberNameIndices[j].Offset], matchedDef.MemberNameIndices[j].Size};
                if (WriteDestinations[j] == UINT64_MAX)
                {
                    // Unknown member name or type changed
                    serializer.Skip(childDef.Size);
                    continue;
                }

)END";
    constexpr char LoadFuncBodyTypeUpgradeMember3[] = R"END(                if (bx::strCmp(memberName, "%s") == 0)
                {
                    auto* fieldPtr = reinterpret_cast<%s*>(objBasePtr + WriteDestinations[j]);
                    isOk = Load(fieldPtr, %u, serializer) && isOk;
                    continue;
                }
)END";
    constexpr char LoadFuncBodyTypeUpgradeEnd[] = R"END(                assert(false);
            }
        }
        assert(isOk);
        return isOk;
    }
)END";
    constexpr char LoadFuncBodyNative[] = R"END(        bool isOk = true;
        for (int32_t i = 0; i < count; ++i)
        {
            isOk = serializer.Read(&obj[i], sizeof(obj[i])) && isOk;
        }
        return isOk;
    }
)END";
    constexpr char LoadFuncBodyEnum2[] = R"END(        bool isOk = true;
        for (int32_t i = 0; i < count; ++i)
        {
            %s& val = (%s&)obj[i];
            isOk = Load(&val, 1, serializer) && isOk;
        }
        return isOk;
    }
)END";
    constexpr char MetadataStart1[] = R"END(
    namespace Meta {
    struct StrRef
    {
        uint16_t Offset;
        uint16_t Size;
    };
    
    struct TypeDef
    {
        uint32_t Size = 0;
        uint32_t Hash = 0;
        char Name[64]{"Dummy"};
        uint16_t ChildCount = 0;
        uint16_t ChildIndices[64]{0};
        StrRef MemberNameIndices[64]{0};
    };

    struct MetadataTable
    {
        TypeDef TypeDefinitions[%u]
        {
)END";

    constexpr char MetadataTypeEntry6[] = R"END(            TypeDef{sizeof(%s), %u, "%s", %u, {%s}, {%s}},
)END";
    constexpr char MetadataEnd1[] = R"END(        };
        char MemberNameBuffer[64*64*64]{"%s"};
    };

    constexpr MetadataTable Metadata;
    }
)END";
} // namespace WriteTemplates

void CppFileWriter::InitBuffer(WriteBuffer& buf)
{
    if (buf.Data != nullptr)
    {
        LOG_ERROR(0, "Multiple writes not supported yet!");
        return;
    }
    buf.Data =
        reinterpret_cast<char*>(VirtualAlloc(nullptr, BufferRequestSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE));
    if (buf.Data == nullptr)
    {
        LOG_ERROR(0, "Failed to allocate write memory!");
        return;
    }
    buf.WrittenBytes = 0;
}

bool CppFileWriter::WriteBufferToDisk(const bx::FilePath& path, const WriteBuffer& buf)
{
    bx::Error error;
    bx::FileWriter writer;

    LOG(0, "Writing to %s", path.getCPtr());
    if (!writer.open(path, false, &error))
    {
        LOG_ERROR(0, "Failed to open output file: %s", error.getMessage().getCPtr());
        return false;
    }

    writer.write(buf.Data, buf.WrittenBytes, &error);
    if (!error.isOk())
    {
        LOG_ERROR(0, "Failed to write to output file: %s", error.getMessage().getCPtr());
    }
    writer.close();
    LOG(0, "Finished writing!");
    return true;
}

void CppFileWriter::WriteInternal(WriteBuffer& buf, const char* templateStr, va_list args)
{
    if (buf.Data == nullptr)
    {
        LOG_ERROR(0, "Wrote too early!");
        return;
    }

    buf.WrittenBytes += bx::vsnprintf(&buf.Data[buf.WrittenBytes], BufferRequestSize, templateStr, args);
}

void CppFileWriter::PrintTypeName(
    char* buf, int32_t bufSize, Def::TypeRef type, const Def::DefinitionFile& definitions, bool printActualEnumType)
{
    if (buf == nullptr || bufSize == 0 || !IsValid(type))
    {
        LOG_ERROR(0, "Invalid PrintTypeName call!");
        return;
    }
    if (type.FieldKind == Def::EFieldType::DefinedClass)
    {
        auto& t = definitions.Types[type.TypeIdx];
        if ((uint32_t)t.TypeFlags & (uint32_t)Def::ETypeFlags::IsNative)
        {
            bx::strCopy(buf, bufSize, t.NativeCName);
        }
        else
        {
            bx::strCopy(buf, bufSize, t.Name);
        }
    }
    else if (type.FieldKind == Def::EFieldType::DefinedEnum)
    {
        bx::strCopy(buf, bufSize, definitions.Enums[type.TypeIdx].Name);
        if (printActualEnumType)
        {
            bx::strCat(buf, bufSize, "::Enum");
        }
    }
}

void CppFileWriter::WriteEnums(const Def::DefinitionFile& definitions)
{
    for (int32_t enumIdx = 0; enumIdx < definitions.EnumCount; ++enumIdx)
    {
        const Def::Enum& e = definitions.Enums[enumIdx];

        if (!IsValid(e.EnumType))
        {
            LOG_ERROR(0, "Invalid enum type (enum %i)", enumIdx);
            continue;
        }

        char Buf[Def::MaxNameLength]{0};
        PrintTypeName(Buf, sizeof(Buf), e.EnumType, definitions);
        Write(WriteTemplates::EnumHeader4, e.Name, definitions.TypeCount + enumIdx, e.EntryCount, Buf);
        for (int32_t entryIdx = 0; entryIdx < e.EntryCount; ++entryIdx)
        {
            Write(WriteTemplates::EnumField1, e.EntryNames[entryIdx]);
        }

        Write(WriteTemplates::EnumNamesStart2, "EntryNames", Def::MaxNameLength);
        for (int32_t entryIdx = 0; entryIdx < e.EntryCount; ++entryIdx)
        {
            Write(WriteTemplates::EnumNamesEntry1, e.EntryNames[entryIdx]);
        }

        for (int32_t extraIdx = 0; extraIdx < e.ExtraStringFieldCount; ++extraIdx)
        {
            Write(WriteTemplates::EnumNamesStart2, e.ExtraStringFieldNames[extraIdx], Def::MaxNameLength);
            for (int32_t entryIdx = 0; entryIdx < e.EntryCount; ++entryIdx)
            {
                Write(WriteTemplates::EnumNamesEntry1, e.ExtraStringFields[entryIdx][extraIdx]);
            }
        }

        Write(WriteTemplates::EnumNamesEnd);
    }
}
void CppFileWriter::WriteTypes(const Def::DefinitionFile& definitions)
{
    for (int32_t typeIdx = 0; typeIdx < definitions.TypeCount; ++typeIdx)
    {
        const Def::Type& t = definitions.Types[typeIdx];
        if ((uint32_t)t.TypeFlags & (uint32_t)Def::ETypeFlags::IsNative) continue;

        Write(WriteTemplates::StructHeader2, t.Name, typeIdx);
        for (int32_t fieldIdx = 0; fieldIdx < t.FieldCount; ++fieldIdx)
        {
            char Type[64]{0};
            PrintTypeName(Type, sizeof(Type), t.FieldTypes[fieldIdx], definitions);
            char Array[32]{0};
            uint32_t ArraySize = t.FieldArraySizes[fieldIdx];
            if (ArraySize > 0 && ArraySize != UINT32_MAX)
            {
                bx::snprintf(Array, sizeof(Array), "[%u]", ArraySize);
            }
            Write(WriteTemplates::StructField4, Type, t.FieldNames[fieldIdx], Array, t.FieldValues[fieldIdx]);
        }
        Write(WriteTemplates::StructEnd);
    }
}
void CppFileWriter::WriteSaveLoadMethods(const Def::DefinitionFile& definitions)
{
    char nameBuf[64]{0};
    char fieldBuf[64]{0};
    for (int32_t enumIdx = 0; enumIdx < definitions.EnumCount; ++enumIdx)
    {
        const Def::Enum& e = definitions.Enums[enumIdx];
        if (!IsValid(e.EnumType))
        {
            LOG_ERROR(0, "Invalid enum!");
            continue;
        }
        bx::snprintf(nameBuf, sizeof(nameBuf), "%s::Enum", e.Name);
        PrintTypeName(fieldBuf, sizeof(fieldBuf), e.EnumType, definitions);

        Write(WriteTemplates::SaveFuncHeader1, nameBuf);
        Write(WriteTemplates::LoadFuncHeader1, nameBuf);

        WriteCpp(WriteTemplates::SaveFuncBodyStart1, nameBuf);
        WriteCpp(WriteTemplates::SaveFuncBodyEnum1, fieldBuf);
        WriteCpp(WriteTemplates::SaveFuncBodyEnd);

        WriteCpp(WriteTemplates::LoadFuncBodyStart1, nameBuf);
        WriteCpp(WriteTemplates::LoadFuncBodyEnum2, fieldBuf, fieldBuf);
    }
    for (uint16_t typeIdx = 0; typeIdx < definitions.TypeCount; ++typeIdx)
    {
        const Def::Type& t = definitions.Types[typeIdx];
        char typeName[Def::MaxNameLength]{0};
        PrintTypeName(typeName, sizeof(typeName), {typeIdx, Def::EFieldType::DefinedClass}, definitions);

        Write(WriteTemplates::SaveFuncHeader1, typeName);
        Write(WriteTemplates::LoadFuncHeader1, typeName);

        WriteCpp(WriteTemplates::SaveFuncBodyStart1, typeName);
        if ((int32_t)t.TypeFlags & (int32_t)Def::ETypeFlags::IsNative)
        {
            WriteCpp(WriteTemplates::SaveFuncBodyNative);
        }
        else
        {
            for (int32_t fieldIdx = 0; fieldIdx < t.FieldCount; ++fieldIdx)
            {
                WriteCpp(WriteTemplates::SaveFuncBodyType3,
                         t.FieldArraySizes[fieldIdx] > 0 ? "" : "&",
                         t.FieldNames[fieldIdx],
                         bx::max(1, t.FieldArraySizes[fieldIdx]));
            }
        }
        WriteCpp(WriteTemplates::SaveFuncBodyEnd);

        WriteCpp(WriteTemplates::LoadFuncBodyStart1, typeName);
        if ((int32_t)t.TypeFlags & (int32_t)Def::ETypeFlags::IsNative)
        {
            WriteCpp(WriteTemplates::LoadFuncBodyNative);
        }
        else
        {
            WriteCpp(WriteTemplates::LoadFuncBodySetupStart2, typeName, typeName);
            for (int32_t fieldIdx = 0; fieldIdx < t.FieldCount; ++fieldIdx)
            {
                WriteCpp(WriteTemplates::LoadFuncBodyQuickLoad2,
                         t.FieldArraySizes[fieldIdx] > 0 ? "" : "&",
                         t.FieldNames[fieldIdx],
                         bx::max(1, t.FieldArraySizes[fieldIdx]));
            }
            WriteCpp(WriteTemplates::LoadFuncBodySetupEnd);
            char fieldTypeName[Def::MaxNameLength]{0};
            for (int32_t fieldIdx = 0; fieldIdx < t.FieldCount; ++fieldIdx)
            {
                const char* fieldName = t.FieldNames[fieldIdx];
                PrintTypeName(fieldTypeName, sizeof(fieldTypeName), t.FieldTypes[fieldIdx], definitions, false);
                WriteCpp(WriteTemplates::LoadFuncBodyMemberCheck4, fieldName, fieldTypeName, typeName, fieldName);
            }
            WriteCpp(WriteTemplates::LoadFuncBodyTypeUpgradeStart);
            for (uint32_t fieldIdx = 0; fieldIdx < t.FieldCount; ++fieldIdx)
            {
                const char* fieldName = t.FieldNames[fieldIdx];
                PrintTypeName(fieldTypeName, sizeof(fieldTypeName), t.FieldTypes[fieldIdx], definitions);
                WriteCpp(WriteTemplates::LoadFuncBodyTypeUpgradeMember3,
                         fieldName,
                         fieldTypeName,
                         bx::max(1, t.FieldArraySizes[fieldIdx]));
            }
            WriteCpp(WriteTemplates::LoadFuncBodyTypeUpgradeEnd);
        }
    }
}

namespace
{
    char MemberNameBuffer[64 * 64 * 64]{0};
}
void CppFileWriter::WriteMetadata(const Def::DefinitionFile& definitions)
{
    uint32_t memberNameBufferIdx = 0;
    Write(WriteTemplates::MetadataStart1, definitions.TypeCount + definitions.EnumCount);
    for (uint16_t i = 0; i < definitions.TypeCount; ++i)
    {
        auto& type = definitions.Types[i];

        char fieldStr[256]{0};
        for (int32_t j = 0; j < type.FieldCount; ++j)
        {
            if (j != 0) bx::strCat(fieldStr, sizeof(fieldStr), ", ");
            char numBuf[8]{0};
            int32_t idx = type.FieldTypes[j].TypeIdx;
            if (type.FieldTypes[j].FieldKind == Def::EFieldType::DefinedEnum) idx += definitions.TypeCount;
            bx::snprintf(numBuf, sizeof(numBuf), "%u", idx);
            bx::strCat(fieldStr, sizeof(fieldStr), numBuf);
        }

        char typeStr[Def::MaxNameLength]{0};
        PrintTypeName(typeStr, sizeof(typeStr), {i, Def::EFieldType::DefinedClass}, definitions);

        char memberIdxStr[256]{0};
        for (int32_t j = 0; j < type.FieldCount; ++j)
        {
            if (j != 0) bx::strCat(memberIdxStr, sizeof(memberIdxStr), ", ");
            char numBuf[16]{0};
            bx::snprintf(numBuf, sizeof(numBuf), "{%u, %u}", memberNameBufferIdx, bx::strLen(type.FieldNames[j]));
            bx::strCat(memberIdxStr, sizeof(memberIdxStr), numBuf);

            memberNameBufferIdx += bx::strCopy(&MemberNameBuffer[memberNameBufferIdx],
                                               sizeof(MemberNameBuffer) - memberNameBufferIdx,
                                               type.FieldNames[j]);
        }
        Write(
            WriteTemplates::MetadataTypeEntry6, typeStr, type.Hash, type.Name, type.FieldCount, fieldStr, memberIdxStr);
    }
    for (uint16_t i = 0; i < definitions.EnumCount; ++i)
    {
        auto& enumType = definitions.Enums[i];

        char typeStr[Def::MaxNameLength]{0};
        PrintTypeName(typeStr, sizeof(typeStr), {i, Def::EFieldType::DefinedEnum}, definitions);

        Write(WriteTemplates::MetadataTypeEntry6, typeStr, enumType.Hash, enumType.Name, 0, "", "");
    }
    Write(WriteTemplates::MetadataEnd1, MemberNameBuffer);
}

void CppFileWriter::GenerateCpp(const bx::FilePath& outDir, const Def::DefinitionFile& definitions)
{
    LOG(0, "Generating output files...");

    InitBuffer(HeaderWrite);
    InitBuffer(CppWrite);

    Write(WriteTemplates::FileHeaderStart);
    WriteCpp(WriteTemplates::FileCppStart);
    WriteEnums(definitions);
    WriteTypes(definitions);
    WriteSaveLoadMethods(definitions);
    WriteMetadata(definitions);
    Write(WriteTemplates::FileEnd);
    WriteCpp(WriteTemplates::FileEnd);

    // std::filesystem::create_directory(outDir.getCPtr());
    bx::FilePath writePathHeader = outDir;
    bx::FilePath writePathCpp = outDir;
    writePathHeader.join("Generated.h");
    writePathCpp.join("Generated.cpp");
    WriteBufferToDisk(writePathHeader, HeaderWrite);
    WriteBufferToDisk(writePathCpp, CppWrite);
}

void CppFileWriter::Write(const char* templateStr, ...)
{
    va_list args;
    va_start(args, templateStr);
    WriteInternal(HeaderWrite, templateStr, args);
    va_end(args);
}

void CppFileWriter::WriteCpp(const char* templateStr, ...)
{
    va_list args;
    va_start(args, templateStr);
    WriteInternal(CppWrite, templateStr, args);
    va_end(args);
}
