#include "wmo_loader.h"
#include "wow_chunk_reader.h"

#include <godot_cpp/core/class_db.hpp>
#include <godot_cpp/variant/color.hpp>
#include <godot_cpp/variant/utility_functions.hpp>
#include <godot_cpp/variant/packed_string_array.hpp>
#include <godot_cpp/variant/quaternion.hpp>

#include <fstream>
#include <filesystem>
#include <cstring>
#include <cstdio>

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

// ─────────────────────────────────────────────────────────────────────────────
// Raw WMO binary structures (packed, no padding)
// ─────────────────────────────────────────────────────────────────────────────
#pragma pack(push, 1)

struct MOHDChunk {          // Root header
    uint32_t nMaterials;
    uint32_t nGroups;
    uint32_t nPortals;
    uint32_t nLights;
    uint32_t nModels;
    uint32_t nDoodads;
    uint32_t nSets;
    uint32_t ambColor;
    uint32_t areaID;
    float    bboxMin[3];
    float    bboxMax[3];
    uint16_t flags;
    uint16_t numLiquids;
};

struct MOMTEntry {          // Material
    uint32_t flags;
    uint32_t shader;
    uint32_t blendMode;
    uint32_t texUnit0;      // index into MOTX
    uint32_t sidnColor;
    uint32_t frameSidnColor;
    uint32_t texUnit1;
    uint32_t diffColor;
    uint32_t groundType;
    uint32_t texUnit2;
    uint32_t color2;
    uint32_t flags2;
    uint32_t runTimeData[4];
};

struct MOGPHeader {         // Group header (inside MOGP chunk)
    uint32_t groupNameOfs;
    uint32_t descNameOfs;
    uint32_t flags;
    float    bboxMin[3];
    float    bboxMax[3];
    uint16_t portalStart;
    uint16_t portalCount;
    uint16_t transBatchCount;
    uint16_t intBatchCount;
    uint16_t extBatchCount;
    uint16_t padding;
    uint8_t  fogIds[4];
    uint32_t liquidType;
    uint32_t groupID;
    uint32_t unk32a;
    uint32_t unk32b;
};

struct MOBAEntry {          // Render batch
    int16_t  bboxMin[3];
    int16_t  bboxMax[3];
    uint32_t indexStart;
    uint16_t indexCount;
    uint16_t vertexStart;
    uint16_t vertexEnd;
    uint8_t  flags;
    uint8_t  materialId;
};

struct MOVTEntry { float x, y, z; };
struct MONREntry { float x, y, z; };
struct MOTVEntry { float u, v; };
struct MOCVEntry { uint8_t b, g, r, a; };

struct MLIQHeader {         // Liquid header inside MLIQ chunk
    uint32_t xverts;        // = xtiles + 1
    uint32_t yverts;        // = ytiles + 1
    uint32_t xtiles;
    uint32_t ytiles;
    float    corner[3];     // base position (WMO local, Z-up)
    uint16_t materialId;
};

struct MLIQVertex {         // 8 bytes per vertex (interpretation depends on liquid type)
    uint32_t data;          // flow info or s,t for magma — we ignore
    float    height;
};

struct MODSEntry {          // Doodad set: 32 bytes
    char     name[20];
    uint32_t startDoodad;   // first MODD index
    uint32_t nDoodads;
    uint32_t unused;
};

struct MODDEntry {          // Doodad placement: 40 bytes
    uint32_t nameOfsAndFlags;   // low 24 = MODN byte offset, high 8 = flags
    float    pos[3];            // WMO local (Z-up)
    float    rot[4];            // quaternion (qx, qy, qz, qw) in WMO local
    float    scale;
    uint8_t  color[4];          // BGRA tint
};

#pragma pack(pop)

// MLIQ tile size in WoW yards (same constant as ADT MCNK unit cells).
static constexpr float WMO_LIQUID_UNIT_SIZE = 4.1666666f;

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────
std::vector<uint8_t> WMOLoader::read_file(const std::string &path) {
    std::ifstream f(path, std::ios::binary | std::ios::ate);
    if (!f) return {};
    auto sz = f.tellg();
    f.seekg(0);
    std::vector<uint8_t> buf(sz);
    f.read(reinterpret_cast<char *>(buf.data()), sz);
    return buf;
}

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

String WMOLoader::to_godot(const std::string &s) {
    return String(s.c_str());
}

static Color color_from_bgra_u32(uint32_t value) {
    const float b = (value & 0xFFu) / 255.0f;
    const float g = ((value >> 8) & 0xFFu) / 255.0f;
    const float r = ((value >> 16) & 0xFFu) / 255.0f;
    const float a = ((value >> 24) & 0xFFu) / 255.0f;
    return Color(r, g, b, a);
}

static Color color_from_rgba_u32(uint32_t value) {
    const float r = (value & 0xFFu) / 255.0f;
    const float g = ((value >> 8) & 0xFFu) / 255.0f;
    const float b = ((value >> 16) & 0xFFu) / 255.0f;
    const float a = ((value >> 24) & 0xFFu) / 255.0f;
    return Color(r, g, b, a);
}

// ─────────────────────────────────────────────────────────────────────────────
// Root parser
// ─────────────────────────────────────────────────────────────────────────────
Dictionary WMOLoader::parse_root(const std::vector<uint8_t> &buf) {
    Dictionary root;
    if (buf.empty()) return root;

    ChunkReader reader(buf.data(), buf.size());
    WoWChunk chunk;

    PackedStringArray textures;
    Array materials;
    std::vector<std::string> tex_strings;

    // Two-pass: first collect raw chunks, then build (MODD/MODN/MODS may appear in any order).
    const uint8_t *modn_data = nullptr;
    size_t         modn_size = 0;
    const MODSEntry *mods = nullptr;
    uint32_t         mods_count = 0;
    const MODDEntry *modd = nullptr;
    uint32_t         modd_count = 0;

    while (reader.next(chunk)) {

        if (chunk.is("MOTX")) {
            // Null-separated texture paths
            auto names = parse_string_block(chunk);
            tex_strings = names;
            for (auto &n : names)
                textures.push_back(to_godot(n));

        } else if (chunk.is("MOMT")) {
            uint32_t n = chunk.count_of(sizeof(MOMTEntry));
            auto *mt   = chunk.array<MOMTEntry>();
            for (uint32_t i = 0; i < n; ++i) {
                Dictionary mat;
                mat["flags"]      = (int)mt[i].flags;
                mat["shader"]     = (int)mt[i].shader;
                mat["blend_mode"] = (int)mt[i].blendMode;
                mat["emissive_color"] = color_from_bgra_u32(mt[i].sidnColor);
                mat["diffuse_color"] = color_from_rgba_u32(mt[i].diffColor);
                mat["color2"] = color_from_rgba_u32(mt[i].color2);
                mat["flags2"] = (int)mt[i].flags2;
                // texUnit offsets → find index in tex_strings
                auto tex_idx = [&](uint32_t ofs) -> int {
                    uint32_t cur = 0;
                    for (int j = 0; j < (int)tex_strings.size(); ++j) {
                        if (cur == ofs) return j;
                        cur += (uint32_t)tex_strings[j].size() + 1;
                    }
                    return -1;
                };
                mat["texture0"] = tex_idx(mt[i].texUnit0);
                mat["texture1"] = tex_idx(mt[i].texUnit1);
                mat["texture2"] = tex_idx(mt[i].texUnit2);
                materials.push_back(mat);
            }
        } else if (chunk.is("MODS")) {
            mods = chunk.array<MODSEntry>();
            mods_count = chunk.count_of(sizeof(MODSEntry));
        } else if (chunk.is("MODN")) {
            modn_data = chunk.data;
            modn_size = chunk.size;
        } else if (chunk.is("MODD")) {
            modd = chunk.array<MODDEntry>();
            modd_count = chunk.count_of(sizeof(MODDEntry));
        }
    }

    root["textures"]  = textures;
    root["materials"] = materials;

    // ── Doodad sets ─────────────────────────────────────────────────────────
    Array doodad_sets;
    for (uint32_t i = 0; i < mods_count; ++i) {
        Dictionary s;
        char buf[21]; std::memcpy(buf, mods[i].name, 20); buf[20] = 0;
        s["name"]  = String(buf);
        s["start"] = (int)mods[i].startDoodad;
        s["count"] = (int)mods[i].nDoodads;
        doodad_sets.push_back(s);
    }
    root["doodad_sets"] = doodad_sets;

    // ── Doodad placements ───────────────────────────────────────────────────
    // Position/rotation come in WMO local (Z-up). Convert to Godot Y-up the
    // same way as MOVT vertices: (lx, ly, lz) → (-ly, lz, -lx). Quaternion
    // vector part transforms as a vector under that basis change; the scalar
    // part is kept (signed `qw` → `qw`). The basis change has det = -1 which
    // reverses winding; for rotations of M2 prototypes that already live in
    // Godot-space (built via M2Builder), this gives the visually correct
    // orientation.
    Array doodad_placements;
    for (uint32_t i = 0; i < modd_count; ++i) {
        const MODDEntry &d = modd[i];

        // Resolve doodad name from MODN by byte offset.
        uint32_t name_ofs = d.nameOfsAndFlags & 0x00FFFFFFu;
        String   name;
        if (modn_data && name_ofs < modn_size) {
            const char *s = reinterpret_cast<const char *>(modn_data + name_ofs);
            // Bound by remaining MODN bytes.
            size_t max_len = modn_size - name_ofs;
            size_t len = 0;
            while (len < max_len && s[len] != 0) ++len;
            name = String(std::string(s, len).c_str());
        }

        Dictionary p;
        p["name"]  = name;
        p["flags"] = (int)((d.nameOfsAndFlags >> 24) & 0xFFu);
        p["pos"]   = Vector3(-d.pos[1], d.pos[2], -d.pos[0]);
        p["rot"]   = Quaternion(-d.rot[1], d.rot[2], -d.rot[0], d.rot[3]);
        p["scale"] = d.scale;
        p["color"] = Color(d.color[2] / 255.f, d.color[1] / 255.f,
                           d.color[0] / 255.f, d.color[3] / 255.f);
        doodad_placements.push_back(p);
    }
    root["doodad_placements"] = doodad_placements;

    return root;
}

// ─────────────────────────────────────────────────────────────────────────────
// Group parser
// ─────────────────────────────────────────────────────────────────────────────
Dictionary WMOLoader::parse_group(const std::vector<uint8_t> &buf) {
    Dictionary group;
    if (buf.empty()) return group;

    ChunkReader top(buf.data(), buf.size());
    WoWChunk chunk;

    // Skip MVER, find MOGP
    while (top.next(chunk)) {
        if (!chunk.is("MOGP")) continue;

        // MOGP data = MOGPHeader + sub-chunks
        if (chunk.size < sizeof(MOGPHeader)) break;
        // const auto &hdr = *reinterpret_cast<const MOGPHeader *>(chunk.data);

        // Sub-chunks start after the header
        size_t sub_ofs  = sizeof(MOGPHeader);
        ChunkReader sub(chunk.data + sub_ofs, chunk.size - sub_ofs);
        WoWChunk sc;

        PackedVector3Array vertices, normals;
        PackedVector2Array uvs;
        PackedColorArray   colors;
        PackedInt32Array   indices;
        Array              batches;

        while (sub.next(sc)) {
            if (sc.is("MOVT")) {
                uint32_t n = sc.count_of(sizeof(MOVTEntry));
                auto *v    = sc.array<MOVTEntry>();
                vertices.resize(n);
                for (uint32_t i = 0; i < n; ++i) {
                    float gx, gy, gz;
                    // WMO local coords → negate Z for Godot Y-up
                    gx = -v[i].y;
                    gy =  v[i].z;
                    gz = -v[i].x;
                    vertices[i] = Vector3(gx, gy, gz);
                }
            } else if (sc.is("MONR")) {
                uint32_t n = sc.count_of(sizeof(MONREntry));
                auto *v    = sc.array<MONREntry>();
                normals.resize(n);
                for (uint32_t i = 0; i < n; ++i)
                    normals[i] = Vector3(-v[i].y, v[i].z, -v[i].x);

            } else if (sc.is("MOTV")) {
                uint32_t n = sc.count_of(sizeof(MOTVEntry));
                auto *v    = sc.array<MOTVEntry>();
                uvs.resize(n);
                for (uint32_t i = 0; i < n; ++i)
                    uvs[i] = Vector2(v[i].u, v[i].v);

            } else if (sc.is("MOVI")) {
                uint32_t n = sc.size / 2;
                auto *idx  = sc.array<uint16_t>();
                indices.resize(n);
                for (uint32_t i = 0; i < n; ++i)
                    indices[i] = (int)idx[i];

            } else if (sc.is("MOCV")) {
                uint32_t n = sc.count_of(sizeof(MOCVEntry));
                auto *v    = sc.array<MOCVEntry>();
                colors.resize(n);
                for (uint32_t i = 0; i < n; ++i)
                    colors[i] = Color(v[i].r / 255.f, v[i].g / 255.f,
                                      v[i].b / 255.f, v[i].a / 255.f);

            } else if (sc.is("MOBA")) {
                uint32_t n = sc.count_of(sizeof(MOBAEntry));
                auto *ba   = sc.array<MOBAEntry>();
                for (uint32_t i = 0; i < n; ++i) {
                    Dictionary b;
                    b["index_start"]  = (int)ba[i].indexStart;
                    b["index_count"]  = (int)ba[i].indexCount;
                    b["vertex_start"] = (int)ba[i].vertexStart;
                    b["vertex_count"] = (int)ba[i].vertexEnd - (int)ba[i].vertexStart + 1;
                    b["material_id"]  = (int)ba[i].materialId;
                    b["flags"]        = (int)ba[i].flags;
                    batches.push_back(b);
                }
            } else if (sc.is("MLIQ") && sc.size >= sizeof(MLIQHeader)) {
                const auto &h = *reinterpret_cast<const MLIQHeader *>(sc.data);

                size_t need = sizeof(MLIQHeader)
                            + (size_t)h.xverts * h.yverts * sizeof(MLIQVertex)
                            + (size_t)h.xtiles * h.ytiles;
                if (need <= sc.size && h.xverts > 0 && h.yverts > 0
                                    && h.xtiles > 0 && h.ytiles > 0) {
                    auto *verts = reinterpret_cast<const MLIQVertex *>(
                        sc.data + sizeof(MLIQHeader));
                    auto *tiles = sc.data + sizeof(MLIQHeader)
                                + (size_t)h.xverts * h.yverts * sizeof(MLIQVertex);

                    PackedFloat32Array heights;
                    heights.resize((int)(h.xverts * h.yverts));
                    {
                        float *dst = heights.ptrw();
                        for (uint32_t i = 0; i < h.xverts * h.yverts; ++i)
                            dst[i] = verts[i].height;
                    }

                    PackedByteArray tile_flags;
                    tile_flags.resize((int)(h.xtiles * h.ytiles));
                    std::memcpy(tile_flags.ptrw(), tiles,
                                (size_t)h.xtiles * h.ytiles);

                    Dictionary liquid;
                    liquid["xverts"]      = (int)h.xverts;
                    liquid["yverts"]      = (int)h.yverts;
                    liquid["xtiles"]      = (int)h.xtiles;
                    liquid["ytiles"]      = (int)h.ytiles;
                    // Corner in WMO local (Z-up). Builder converts to Godot.
                    liquid["corner"]      = Vector3(h.corner[0], h.corner[1], h.corner[2]);
                    liquid["material_id"] = (int)h.materialId;
                    liquid["unit_size"]   = WMO_LIQUID_UNIT_SIZE;
                    liquid["heights"]     = heights;
                    liquid["tiles"]       = tile_flags;
                    group["liquid"]       = liquid;
                }
            }
        }

        group["vertices"] = vertices;
        group["normals"]  = normals;
        group["uvs"]      = uvs;
        group["colors"]   = colors;
        group["indices"]  = indices;
        group["batches"]  = batches;
        break; // only one MOGP per group file
    }

    return group;
}

// ─────────────────────────────────────────────────────────────────────────────
// Public API
// ─────────────────────────────────────────────────────────────────────────────
Dictionary WMOLoader::load_root(const String &path) {
    auto buf = read_file(to_std(path));
    if (buf.empty()) {
        UtilityFunctions::push_error("WMOLoader: cannot read ", path);
        return Dictionary();
    }
    return parse_root(buf);
}

Dictionary WMOLoader::load_group(const String &path) {
    auto buf = read_file(to_std(path));
    if (buf.empty()) {
        UtilityFunctions::push_error("WMOLoader: cannot read ", path);
        return Dictionary();
    }
    return parse_group(buf);
}

Dictionary WMOLoader::load_wmo(const String &path) {
    std::string spath = to_std(path);
    auto root_buf = read_file(spath);
    if (root_buf.empty()) {
        UtilityFunctions::push_error("WMOLoader: cannot read root ", path);
        return Dictionary();
    }

    Dictionary result = parse_root(root_buf);

    // Count groups from MOHD
    uint32_t nGroups = 0;
    {
        ChunkReader r(root_buf.data(), root_buf.size());
        WoWChunk c;
        if (r.find("MOHD", c) && c.size >= sizeof(MOHDChunk))
            nGroups = c.as<MOHDChunk>().nGroups;
    }

    // Build base path (strip .wmo extension)
    std::string base = spath;
    if (base.size() > 4 && base.substr(base.size() - 4) == ".wmo")
        base = base.substr(0, base.size() - 4);

    Array groups;
    for (uint32_t i = 0; i < nGroups; ++i) {
        char suffix[16];
        std::snprintf(suffix, sizeof(suffix), "_%03u.wmo", i);
        std::string gpath = base + suffix;

        auto gbuf = read_file(gpath);
        if (gbuf.empty()) {
            UtilityFunctions::push_warning("WMOLoader: group not found: ",
                to_godot(gpath));
            groups.push_back(Dictionary());
            continue;
        }
        groups.push_back(parse_group(gbuf));
    }

    result["groups"] = groups;
    UtilityFunctions::print("WMOLoader: loaded ", (int)nGroups, " groups from ", path);
    return result;
}

Dictionary WMOLoader::load_wmo_from_bytes(const PackedByteArray &root_bytes,
                                           const Array &group_bytes_array)
{
    std::vector<uint8_t> rbuf(root_bytes.size());
    std::memcpy(rbuf.data(), root_bytes.ptr(), root_bytes.size());

    Dictionary result = parse_root(rbuf);
    Array groups;

    for (int i = 0; i < group_bytes_array.size(); ++i) {
        PackedByteArray gb = group_bytes_array[i];
        if (gb.is_empty()) { groups.push_back(Dictionary()); continue; }
        std::vector<uint8_t> gbuf(gb.size());
        std::memcpy(gbuf.data(), gb.ptr(), gb.size());
        groups.push_back(parse_group(gbuf));
    }

    result["groups"] = groups;
    return result;
}

// ─────────────────────────────────────────────────────────────────────────────
void WMOLoader::_bind_methods() {
    ClassDB::bind_method(D_METHOD("load_wmo",            "path"),                         &WMOLoader::load_wmo);
    ClassDB::bind_method(D_METHOD("load_root",           "path"),                         &WMOLoader::load_root);
    ClassDB::bind_method(D_METHOD("load_group",          "path"),                         &WMOLoader::load_group);
    ClassDB::bind_method(D_METHOD("load_wmo_from_bytes", "root_bytes", "group_bytes_array"), &WMOLoader::load_wmo_from_bytes);
}
