#include <XCEngine/Resources/Shader/ShaderCompilationCache.h>

#include <algorithm>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>

namespace XCEngine {
namespace Resources {

namespace fs = std::filesystem;

namespace {

struct ShaderCompilationCacheFileHeader {
    char magic[8] = { 'X', 'C', 'B', 'C', '0', '1', '\0', '\0' };
    Core::uint32 schemaVersion = kShaderCompilationCacheSchemaVersion;
    Core::uint32 backend = 0;
    Core::uint32 format = 0;
    Core::uint32 reserved = 0;
    Core::uint64 payloadSize = 0;
    Core::uint64 compileKeyHigh = 0;
    Core::uint64 compileKeyLow = 0;
};

Containers::String MakeError(const char* message) {
    return Containers::String(message == nullptr ? "" : message);
}

std::string ToStdString(const Containers::String& text) {
    return std::string(text.CStr(), text.Length());
}

Containers::String ToContainersString(const std::string& text) {
    return Containers::String(text.c_str(), text.size());
}

std::vector<std::string> SplitFields(const std::string& line) {
    std::vector<std::string> fields;
    std::stringstream stream(line);
    std::string field;
    while (std::getline(stream, field, '\t')) {
        fields.push_back(field);
    }
    return fields;
}

bool IsExpectedMagic(const ShaderCompilationCacheFileHeader& header) {
    return std::memcmp(header.magic, "XCBC01", 6) == 0 &&
           header.schemaVersion == kShaderCompilationCacheSchemaVersion;
}

} // namespace

void ShaderCompileKey::Normalize() {
    shaderPath = shaderPath.Trim();
    sourceHash = sourceHash.Trim();
    dependencyHash = dependencyHash.Trim();
    passName = passName.Trim();
    entryPoint = entryPoint.Trim();
    profile = profile.Trim();
    compilerName = compilerName.Trim();
    compilerVersion = compilerVersion.Trim();
    optionsSignature = optionsSignature.Trim();

    std::vector<std::string> normalizedKeywords;
    normalizedKeywords.reserve(keywords.Size());
    for (const Containers::String& keyword : keywords) {
        const Containers::String trimmedKeyword = keyword.Trim();
        if (!trimmedKeyword.Empty()) {
            normalizedKeywords.push_back(ToStdString(trimmedKeyword));
        }
    }

    std::sort(normalizedKeywords.begin(), normalizedKeywords.end());
    normalizedKeywords.erase(
        std::unique(normalizedKeywords.begin(), normalizedKeywords.end()),
        normalizedKeywords.end());

    Containers::Array<Containers::String> canonicalKeywords;
    canonicalKeywords.Reserve(normalizedKeywords.size());
    for (const std::string& keyword : normalizedKeywords) {
        canonicalKeywords.PushBack(ToContainersString(keyword));
    }
    keywords = std::move(canonicalKeywords);
}

Containers::String ShaderCompileKey::BuildSignature() const {
    ShaderCompileKey normalizedKey = *this;
    normalizedKey.Normalize();

    Containers::String signature;
    signature += "shader=";
    signature += normalizedKey.shaderPath;
    signature += "\nsourceHash=";
    signature += normalizedKey.sourceHash;
    signature += "\ndependencyHash=";
    signature += normalizedKey.dependencyHash;
    signature += "\npass=";
    signature += normalizedKey.passName;
    signature += "\nentry=";
    signature += normalizedKey.entryPoint;
    signature += "\nprofile=";
    signature += normalizedKey.profile;
    signature += "\ncompiler=";
    signature += normalizedKey.compilerName;
    signature += "\ncompilerVersion=";
    signature += normalizedKey.compilerVersion;
    signature += "\noptions=";
    signature += normalizedKey.optionsSignature;
    signature += "\nstage=";
    signature += ToContainersString(std::to_string(static_cast<Core::uint32>(normalizedKey.stage)));
    signature += "\nsourceLanguage=";
    signature += ToContainersString(std::to_string(static_cast<Core::uint32>(normalizedKey.sourceLanguage)));
    signature += "\nbackend=";
    signature += ToContainersString(std::to_string(static_cast<Core::uint32>(normalizedKey.backend)));

    for (const Containers::String& keyword : normalizedKey.keywords) {
        signature += "\nkeyword=";
        signature += keyword;
    }

    return signature;
}

Containers::String ShaderCompileKey::BuildCacheKey() const {
    return HashStringToAssetGUID(BuildSignature()).ToString();
}

void ShaderCompilationCache::Initialize(const Containers::String& libraryRoot) {
    m_libraryRoot = libraryRoot.Trim();
    m_databasePath = m_libraryRoot + "/shadercache.db";
    LoadDatabase();
}

void ShaderCompilationCache::Shutdown() {
    m_libraryRoot.Clear();
    m_databasePath.Clear();
    m_records.clear();
}

Containers::String ShaderCompilationCache::BuildCacheKey(const ShaderCompileKey& key) const {
    return key.BuildCacheKey();
}

Containers::String ShaderCompilationCache::BuildBackendDirectoryName(ShaderBackend backend) {
    switch (backend) {
    case ShaderBackend::D3D12:
        return Containers::String("D3D12");
    case ShaderBackend::OpenGL:
        return Containers::String("OpenGL");
    case ShaderBackend::Vulkan:
        return Containers::String("Vulkan");
    case ShaderBackend::Generic:
    default:
        return Containers::String("Generic");
    }
}

AssetGUID ShaderCompilationCache::ComputeKeyGuid(const Containers::String& cacheKey) {
    AssetGUID guid;
    if (AssetGUID::TryParse(cacheKey, guid)) {
        return guid;
    }
    return HashStringToAssetGUID(cacheKey);
}

Containers::String ShaderCompilationCache::BuildCacheRelativePath(const ShaderCompileKey& key) const {
    const Containers::String cacheKey = BuildCacheKey(key);
    const Containers::String shard =
        cacheKey.Length() >= 2 ? cacheKey.Substring(0, 2) : Containers::String("00");
    return Containers::String("ShaderCache/") +
           BuildBackendDirectoryName(key.backend) +
           "/" +
           shard +
           "/" +
           cacheKey +
           ".xcbc";
}

Containers::String ShaderCompilationCache::BuildCacheAbsolutePath(const ShaderCompileKey& key) const {
    if (m_libraryRoot.Empty()) {
        return Containers::String();
    }

    return m_libraryRoot + "/" + BuildCacheRelativePath(key);
}

void ShaderCompilationCache::LoadDatabase() {
    m_records.clear();

    std::ifstream input(m_databasePath.CStr());
    if (!input.is_open()) {
        return;
    }

    std::string line;
    while (std::getline(input, line)) {
        if (line.empty() || line[0] == '#') {
            continue;
        }

        const std::vector<std::string> fields = SplitFields(line);
        if (fields.size() < 5) {
            continue;
        }

        ShaderCacheRecord record;
        record.backend = static_cast<ShaderBackend>(std::stoul(fields[1]));
        record.relativePath = ToContainersString(fields[2]);
        record.format = static_cast<ShaderBytecodeFormat>(std::stoul(fields[3]));
        record.payloadSize = static_cast<Core::uint64>(std::stoull(fields[4]));
        m_records[fields[0]] = record;
    }
}

void ShaderCompilationCache::SaveDatabase() const {
    if (m_databasePath.Empty()) {
        return;
    }

    const fs::path dbPath(m_databasePath.CStr());
    std::error_code ec;
    if (!dbPath.parent_path().empty()) {
        fs::create_directories(dbPath.parent_path(), ec);
        if (ec) {
            return;
        }
    }

    std::ofstream output(dbPath, std::ios::out | std::ios::trunc);
    if (!output.is_open()) {
        return;
    }

    output << "# compileKey\tbackend\trelativePath\tformat\tpayloadSize\n";
    for (const auto& [compileKey, record] : m_records) {
        output << compileKey << '\t'
               << static_cast<Core::uint32>(record.backend) << '\t'
               << ToStdString(record.relativePath) << '\t'
               << static_cast<Core::uint32>(record.format) << '\t'
               << record.payloadSize << '\n';
    }
}

bool ShaderCompilationCache::Store(const ShaderCacheEntry& entry,
                                   Containers::String* outErrorMessage) {
    if (!IsInitialized()) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("ShaderCompilationCache is not initialized.");
        }
        return false;
    }

    const Containers::String cacheKey = BuildCacheKey(entry.key);
    const Containers::String relativePath = BuildCacheRelativePath(entry.key);
    const Containers::String absolutePath = BuildCacheAbsolutePath(entry.key);
    if (absolutePath.Empty()) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("ShaderCompilationCache absolute path is empty.");
        }
        return false;
    }

    std::error_code ec;
    const fs::path cachePath(absolutePath.CStr());
    fs::create_directories(cachePath.parent_path(), ec);
    if (ec) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("Failed to create ShaderCompilationCache directory.");
        }
        return false;
    }

    std::ofstream output(cachePath, std::ios::binary | std::ios::trunc);
    if (!output.is_open()) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("Failed to open ShaderCompilationCache output file.");
        }
        return false;
    }

    const AssetGUID keyGuid = ComputeKeyGuid(cacheKey);
    ShaderCompilationCacheFileHeader fileHeader = {};
    fileHeader.backend = static_cast<Core::uint32>(entry.key.backend);
    fileHeader.format = static_cast<Core::uint32>(entry.format);
    fileHeader.payloadSize = static_cast<Core::uint64>(entry.payload.Size());
    fileHeader.compileKeyHigh = keyGuid.high;
    fileHeader.compileKeyLow = keyGuid.low;
    output.write(reinterpret_cast<const char*>(&fileHeader), sizeof(fileHeader));
    if (!output) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("Failed to write ShaderCompilationCache header.");
        }
        return false;
    }

    if (!entry.payload.Empty()) {
        output.write(reinterpret_cast<const char*>(entry.payload.Data()),
                     static_cast<std::streamsize>(entry.payload.Size()));
        if (!output) {
            if (outErrorMessage != nullptr) {
                *outErrorMessage = MakeError("Failed to write ShaderCompilationCache payload.");
            }
            return false;
        }
    }

    ShaderCacheRecord record;
    record.backend = entry.key.backend;
    record.format = entry.format;
    record.relativePath = relativePath;
    record.payloadSize = static_cast<Core::uint64>(entry.payload.Size());
    m_records[ToStdString(cacheKey)] = record;
    SaveDatabase();
    return true;
}

bool ShaderCompilationCache::TryLoad(const ShaderCompileKey& key,
                                     ShaderCacheEntry& outEntry,
                                     Containers::String* outErrorMessage) const {
    outEntry = {};
    if (!IsInitialized()) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("ShaderCompilationCache is not initialized.");
        }
        return false;
    }

    const Containers::String cacheKey = BuildCacheKey(key);
    Containers::String relativePath = BuildCacheRelativePath(key);
    const auto recordIt = m_records.find(ToStdString(cacheKey));
    if (recordIt != m_records.end() && !recordIt->second.relativePath.Empty()) {
        relativePath = recordIt->second.relativePath;
    }

    const Containers::String absolutePath = m_libraryRoot + "/" + relativePath;
    std::ifstream input(absolutePath.CStr(), std::ios::binary);
    if (!input.is_open()) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("ShaderCompilationCache entry does not exist.");
        }
        return false;
    }

    ShaderCompilationCacheFileHeader fileHeader = {};
    input.read(reinterpret_cast<char*>(&fileHeader), sizeof(fileHeader));
    if (!input || !IsExpectedMagic(fileHeader)) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("ShaderCompilationCache header is invalid.");
        }
        return false;
    }

    const AssetGUID expectedKeyGuid = ComputeKeyGuid(cacheKey);
    const AssetGUID actualKeyGuid(fileHeader.compileKeyHigh, fileHeader.compileKeyLow);
    if (expectedKeyGuid != actualKeyGuid ||
        fileHeader.backend != static_cast<Core::uint32>(key.backend)) {
        if (outErrorMessage != nullptr) {
            *outErrorMessage = MakeError("ShaderCompilationCache entry does not match the requested key.");
        }
        return false;
    }

    Containers::Array<Core::uint8> payload;
    payload.ResizeUninitialized(static_cast<size_t>(fileHeader.payloadSize));
    if (fileHeader.payloadSize > 0) {
        input.read(reinterpret_cast<char*>(payload.Data()),
                   static_cast<std::streamsize>(fileHeader.payloadSize));
        if (!input) {
            if (outErrorMessage != nullptr) {
                *outErrorMessage = MakeError("Failed to read ShaderCompilationCache payload.");
            }
            return false;
        }
    }

    outEntry.key = key;
    outEntry.format = static_cast<ShaderBytecodeFormat>(fileHeader.format);
    outEntry.payload = std::move(payload);
    return true;
}

} // namespace Resources
} // namespace XCEngine