#include "mpq_manager.h"

#include <godot_cpp/core/class_db.hpp>
#include <godot_cpp/variant/utility_functions.hpp>
#include <godot_cpp/classes/dir_access.hpp>

#include <algorithm>
#include <filesystem>
#include <fstream>
#include <unordered_set>

namespace fs = std::filesystem;
using namespace godot;

// ─────────────────────────────────────────────────────────────────────────────
// WoW 3.3.5a (build 12340) MPQ load order
// Source: Blizzard client, TrinityCore/MaNGOS toolchain
//
// Priority value = open order (higher = wins).
// Base archives 1–7, locale archives 10–20 (locale always wins over base).
// ─────────────────────────────────────────────────────────────────────────────
std::vector<std::pair<std::string, int>>
MPQManager::build_archive_list(const std::string &data_dir, const std::string &locale)
{
    // {relative path from data_dir, priority}
    std::vector<std::pair<std::string, int>> list = {
        // ── Base archives (lowest priority) ───────────────────────────────
        { "common.mpq",            1 },
        { "common-2.mpq",          2 },
        { "expansion.mpq",         3 },
        { "lichking.mpq",          4 },
        { "patch.mpq",             5 },
        { "patch-2.mpq",           6 },
        { "patch-3.mpq",           7 },
        // ── Locale archives (higher priority) ─────────────────────────────
        { locale + "/locale-"           + locale + ".mpq", 10 },
        { locale + "/speech-"           + locale + ".mpq", 11 },
        { locale + "/expansion-locale-" + locale + ".mpq", 12 },
        { locale + "/lichking-locale-"  + locale + ".mpq", 13 },
        { locale + "/patch-"            + locale + ".mpq", 14 },
        { locale + "/patch-"            + locale + "-2.mpq", 15 },
        { locale + "/patch-"            + locale + "-3.mpq", 16 },
    };

    // Prepend data_dir to each entry
    for (auto &entry : list) {
        entry.first = data_dir + "/" + entry.first;
    }
    return list;
}

// ─────────────────────────────────────────────────────────────────────────────
MPQManager::MPQManager() = default;

MPQManager::~MPQManager() {
    close();
}

// ─────────────────────────────────────────────────────────────────────────────
int MPQManager::open_client(const String &base_path, const String &locale)
{
    close();

    std::string data_dir = to_std(base_path) + "/Data";
    std::string loc      = to_std(locale);

    auto archive_list = build_archive_list(data_dir, loc);

    int opened = 0;
    for (auto &[path, priority] : archive_list) {
        if (!fs::exists(path)) {
            // Not every archive is required to be present
            continue;
        }
        if (open_archive(String(path.c_str()), priority)) {
            ++opened;
        }
    }

    if (opened == 0) {
        UtilityFunctions::push_error("MPQManager: no archives found in ", base_path);
    } else {
        UtilityFunctions::print("MPQManager: opened ", opened, " archives from ", base_path);
    }
    return opened;
}

// ─────────────────────────────────────────────────────────────────────────────
bool MPQManager::open_archive(const String &path, int priority)
{
    std::string spath = to_std(path);
    HANDLE hMpq = nullptr;

    if (!SFileOpenArchive(spath.c_str(), 0, MPQ_OPEN_READ_ONLY | STREAM_FLAG_READ_ONLY, &hMpq)) {
        UtilityFunctions::push_warning("MPQManager: failed to open ", path,
            " (error ", (int)GetLastError(), ")");
        return false;
    }

    MPQArchiveEntry entry;
    entry.handle   = hMpq;
    entry.path     = path;
    entry.priority = priority;

    // Insert sorted ascending by priority so index 0 = lowest
    auto it = std::lower_bound(
        m_archives.begin(), m_archives.end(), priority,
        [](const MPQArchiveEntry &e, int p){ return e.priority < p; });
    m_archives.insert(it, entry);

    return true;
}

// ─────────────────────────────────────────────────────────────────────────────
void MPQManager::close()
{
    for (auto &entry : m_archives) {
        if (entry.handle) {
            SFileCloseArchive(entry.handle);
            entry.handle = nullptr;
        }
    }
    m_archives.clear();
}

// ─────────────────────────────────────────────────────────────────────────────
// Internal: search from highest priority (back of vector) to lowest
HANDLE MPQManager::find_file(const std::string &internal_path) const
{
    for (auto it = m_archives.rbegin(); it != m_archives.rend(); ++it) {
        HANDLE hFile = nullptr;
        if (SFileOpenFileEx(it->handle, internal_path.c_str(), 0, &hFile)) {
            SFileCloseFile(hFile);
            return it->handle;
        }
    }
    return nullptr;
}

// ─────────────────────────────────────────────────────────────────────────────
bool MPQManager::has_file(const String &internal_path) const
{
    return find_file(to_std(internal_path)) != nullptr;
}

// ─────────────────────────────────────────────────────────────────────────────
PackedByteArray MPQManager::read_file(const String &internal_path) const
{
    std::string spath = to_std(internal_path);
    HANDLE hArchive = find_file(spath);
    if (!hArchive) {
        UtilityFunctions::push_error("MPQManager: file not found: ", internal_path);
        return PackedByteArray();
    }

    HANDLE hFile = nullptr;
    if (!SFileOpenFileEx(hArchive, spath.c_str(), 0, &hFile)) {
        return PackedByteArray();
    }

    DWORD file_size = SFileGetFileSize(hFile, nullptr);
    if (file_size == SFILE_INVALID_SIZE || file_size == 0) {
        SFileCloseFile(hFile);
        return PackedByteArray();
    }

    PackedByteArray result;
    result.resize((int)file_size);
    uint8_t *ptr = result.ptrw();

    DWORD read = 0;
    bool ok = SFileReadFile(hFile, ptr, file_size, &read, nullptr);
    SFileCloseFile(hFile);

    if (!ok || read != file_size) {
        UtilityFunctions::push_error("MPQManager: read error on ", internal_path);
        return PackedByteArray();
    }
    return result;
}

// ─────────────────────────────────────────────────────────────────────────────
bool MPQManager::extract_file(const String &internal_path, const String &output_path) const
{
    std::string spath  = to_std(internal_path);
    std::string outpath = to_std(output_path);

    HANDLE hArchive = find_file(spath);
    if (!hArchive) {
        UtilityFunctions::push_error("MPQManager: file not found: ", internal_path);
        return false;
    }

    // Ensure parent directory exists
    fs::path out(outpath);
    fs::create_directories(out.parent_path());

    return SFileExtractFile(hArchive, spath.c_str(), outpath.c_str(), 0);
}

// ─────────────────────────────────────────────────────────────────────────────
int MPQManager::extract_files(const String &filter, const String &output_dir) const
{
    if (m_archives.empty()) return 0;

    std::string out_dir = to_std(output_dir);
    std::string sfilter = to_std(filter);

    // Collect unique file names from all archives (highest priority wins)
    // We iterate archives highest→lowest and track already-seen paths.
    std::unordered_set<std::string> seen;
    int count = 0;

    for (auto it = m_archives.rbegin(); it != m_archives.rend(); ++it) {
        SFILE_FIND_DATA fd;
        HANDLE hFind = SFileFindFirstFile(it->handle, sfilter.c_str(), &fd, nullptr);
        if (!hFind) continue;

        do {
            std::string fname(fd.cFileName);
            if (seen.count(fname)) continue; // already extracted from higher-priority archive
            seen.insert(fname);

            // Build output path mirroring internal path
            fs::path out_path = fs::path(out_dir) / fs::path(fname);
            fs::create_directories(out_path.parent_path());

            if (SFileExtractFile(it->handle, fname.c_str(), out_path.string().c_str(), 0)) {
                ++count;
            } else {
                UtilityFunctions::push_warning("MPQManager: failed to extract ", String(fname.c_str()));
            }
        } while (SFileFindNextFile(hFind, &fd));

        SFileFindClose(hFind);
    }

    UtilityFunctions::print("MPQManager: extracted ", count, " files → ", output_dir);
    return count;
}

// ─────────────────────────────────────────────────────────────────────────────
PackedStringArray MPQManager::list_files(const String &filter) const
{
    std::string sfilter = to_std(filter.is_empty() ? String("*") : filter);
    std::unordered_set<std::string> seen;
    PackedStringArray result;

    // Iterate highest priority first so duplicates are attributed to winner
    for (auto it = m_archives.rbegin(); it != m_archives.rend(); ++it) {
        SFILE_FIND_DATA fd;
        HANDLE hFind = SFileFindFirstFile(it->handle, sfilter.c_str(), &fd, nullptr);
        if (!hFind) continue;

        do {
            std::string fname(fd.cFileName);
            if (!seen.count(fname)) {
                seen.insert(fname);
                result.push_back(String(fname.c_str()));
            }
        } while (SFileFindNextFile(hFind, &fd));

        SFileFindClose(hFind);
    }

    return result;
}

// ─────────────────────────────────────────────────────────────────────────────
int MPQManager::get_archive_count() const {
    return (int)m_archives.size();
}

Array MPQManager::get_archive_info() const {
    Array out;
    for (auto it = m_archives.rbegin(); it != m_archives.rend(); ++it) {
        DWORD file_count = 0;
        SFileGetFileInfo(it->handle, SFileMpqNumberOfFiles, &file_count, sizeof(file_count), nullptr);

        Dictionary d;
        d["path"]       = it->path;
        d["priority"]   = it->priority;
        d["file_count"] = (int)file_count;
        out.push_back(d);
    }
    return out;
}

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────
String MPQManager::to_godot(const std::string &s) {
    return String(s.c_str());
}

std::string MPQManager::to_std(const String &s) {
    return std::string(s.utf8().get_data());
}

// ─────────────────────────────────────────────────────────────────────────────
// GDScript bindings
// ─────────────────────────────────────────────────────────────────────────────
void MPQManager::_bind_methods() {
    ClassDB::bind_method(D_METHOD("open_client", "base_path", "locale"),
        &MPQManager::open_client);
    ClassDB::bind_method(D_METHOD("open_archive", "path", "priority"),
        &MPQManager::open_archive);
    ClassDB::bind_method(D_METHOD("close"),
        &MPQManager::close);
    ClassDB::bind_method(D_METHOD("has_file", "internal_path"),
        &MPQManager::has_file);
    ClassDB::bind_method(D_METHOD("read_file", "internal_path"),
        &MPQManager::read_file);
    ClassDB::bind_method(D_METHOD("extract_file", "internal_path", "output_path"),
        &MPQManager::extract_file);
    ClassDB::bind_method(D_METHOD("extract_files", "filter", "output_dir"),
        &MPQManager::extract_files);
    ClassDB::bind_method(D_METHOD("list_files", "filter"),
        &MPQManager::list_files, DEFVAL("*"));
    ClassDB::bind_method(D_METHOD("get_archive_count"),
        &MPQManager::get_archive_count);
    ClassDB::bind_method(D_METHOD("get_archive_info"),
        &MPQManager::get_archive_info);
}
