#include "hmap/hmap.h"

#include <cstring>
#include <fstream>

namespace hmap {

namespace {

class Writer {
public:
    std::vector<uint8_t> buf;
    void u8(uint8_t v) { buf.push_back(v); }
    void u32(uint32_t v) { for (int i = 0; i < 4; ++i) buf.push_back((uint8_t)((v >> (8 * i)) & 0xff)); }
    void i32(int32_t v) { u32((uint32_t)v); }
    void f32(float f) { uint32_t u; std::memcpy(&u, &f, 4); u32(u); }
    void str(const std::string& s) { u32((uint32_t)s.size()); buf.insert(buf.end(), s.begin(), s.end()); }
    void vec2(const Vec2& v) { f32(v.x); f32(v.y); }
    void vec3(const Vec3& v) { f32(v.x); f32(v.y); f32(v.z); }
    void plane(const Plane& p) { vec3(p.n); f32(p.d); }
    void aabb(const AABB& b) { vec3(b.mn); vec3(b.mx); }
    void uv(const UVInfo& u) { vec2(u.offset); vec2(u.scale); f32(u.rotation); }
    void raw(const char* tag4, const Writer& body) {
        buf.insert(buf.end(), tag4, tag4 + 4);
        u32((uint32_t)body.buf.size());
        buf.insert(buf.end(), body.buf.begin(), body.buf.end());
    }
};

class Reader {
public:
    const uint8_t* p;
    size_t n;
    size_t pos = 0;
    bool ok = true;
    Reader(const uint8_t* data, size_t size) : p(data), n(size) {}
    bool need(size_t k) {
        if (!ok || pos + k > n) { ok = false; return false; }
        return true;
    }
    uint8_t u8() { return need(1) ? p[pos++] : 0; }
    uint32_t u32() {
        if (!need(4)) return 0;
        uint32_t v = (uint32_t)p[pos] | ((uint32_t)p[pos + 1] << 8) | ((uint32_t)p[pos + 2] << 16) | ((uint32_t)p[pos + 3] << 24);
        pos += 4;
        return v;
    }
    int32_t i32() { return (int32_t)u32(); }
    float f32() { uint32_t u = u32(); float f; std::memcpy(&f, &u, 4); return f; }
    std::string str() {
        uint32_t len = u32();
        if (!need(len)) return {};
        std::string s((const char*)p + pos, len);
        pos += len;
        return s;
    }
    Vec2 vec2() { Vec2 v; v.x = f32(); v.y = f32(); return v; }
    Vec3 vec3() { Vec3 v; v.x = f32(); v.y = f32(); v.z = f32(); return v; }
    Plane plane() { Plane pl; pl.n = vec3(); pl.d = f32(); return pl; }
    AABB aabb() { AABB b; b.mn = vec3(); b.mx = vec3(); return b; }
    UVInfo uv() { UVInfo u; u.offset = vec2(); u.scale = vec2(); u.rotation = f32(); return u; }
    // Reads an item count and rejects counts that cannot possibly fit in the remaining bytes.
    uint32_t count(size_t minBytesPerItem) {
        uint32_t c = u32();
        if (ok && (size_t)c * minBytesPerItem > n - pos) ok = false;
        return ok ? c : 0;
    }
};

// ------------------------------------------------------------ writers
void writeMats(Writer& w, const std::vector<Material>& mats) {
    w.u32((uint32_t)mats.size());
    for (const Material& m : mats) {
        w.str(m.name); w.str(m.texture); w.vec2(m.tileSize); w.vec3(m.color); w.u32(m.flags);
    }
}

void writeBrushes(Writer& w, const std::vector<Brush>& bs) {
    w.u32((uint32_t)bs.size());
    for (const Brush& b : bs) {
        w.u32(b.id); w.u32(b.owner); w.u32(b.contents); w.u32(b.flags);
        w.u32((uint32_t)b.faces.size());
        for (const Face& f : b.faces) { w.plane(f.plane); w.i32(f.material); w.uv(f.uv); w.u32(f.flags); }
    }
}

void writeMeshes(Writer& w, const std::vector<MeshObject>& ms) {
    w.u32((uint32_t)ms.size());
    for (const MeshObject& m : ms) {
        w.u32(m.id); w.u32(m.owner); w.u32(m.flags); w.str(m.name);
        w.u32((uint32_t)m.verts.size());
        for (const Vec3& v : m.verts) w.vec3(v);
        w.u32((uint32_t)m.polys.size());
        for (const MeshPoly& p : m.polys) {
            w.i32(p.material); w.u32(p.smooth); w.uv(p.proj);
            w.u32((uint32_t)p.idx.size());
            for (uint32_t i : p.idx) w.u32(i);
            w.u32((uint32_t)p.uv.size());
            for (const Vec2& t : p.uv) w.vec2(t);
        }
    }
}

void writeEntities(Writer& w, const std::vector<Entity>& es) {
    w.u32((uint32_t)es.size());
    for (const Entity& e : es) {
        w.u32(e.id); w.u32(e.flags); w.str(e.classname);
        w.vec3(e.origin); w.vec3(e.angles); w.vec3(e.scale);
        w.u32((uint32_t)e.props.size());
        for (const KeyValue& kv : e.props) { w.str(kv.key); w.str(kv.value); }
    }
}

void writeCollisionBrush(Writer& w, const CollisionBrush& b) {
    w.u32(b.contents); w.u32(b.sourceId); w.aabb(b.bounds);
    w.u32((uint32_t)b.planes.size());
    for (const Plane& p : b.planes) w.plane(p);
}

void writeCompiledChunk(Writer& w, const CompiledMap& cm) {
    w.vec3(cm.ambient); w.f32(cm.unitsPerMeter); w.u8(cm.baked ? 1 : 0); w.aabb(cm.bounds);
    w.u32((uint32_t)cm.batches.size());
    for (const Batch& b : cm.batches) {
        w.i32(b.material); w.u8((uint8_t)b.source); w.i32(b.model);
        w.u32((uint32_t)b.vertices.size());
        for (const Vertex& v : b.vertices) { w.vec3(v.pos); w.vec3(v.normal); w.vec2(v.uv); w.u32(v.color); }
        w.u32((uint32_t)b.indices.size());
        for (uint32_t i : b.indices) w.u32(i);
    }
    w.u32((uint32_t)cm.models.size());
    for (const CompiledModel& m : cm.models) {
        w.u32(m.entityId); w.aabb(m.bounds);
        w.u32((uint32_t)m.brushes.size());
        for (const CollisionBrush& b : m.brushes) writeCollisionBrush(w, b);
    }
    const Bsp& t = cm.bsp;
    w.i32(t.root); w.u32(t.depth); w.aabb(t.bounds);
    w.u32((uint32_t)t.planes.size());
    for (const Plane& p : t.planes) w.plane(p);
    w.u32((uint32_t)t.nodes.size());
    for (const BspNode& n : t.nodes) { w.i32(n.plane); w.i32(n.child[0]); w.i32(n.child[1]); }
    w.u32((uint32_t)t.leaves.size());
    for (const BspLeaf& l : t.leaves) { w.u32(l.firstBrush); w.u32(l.brushCount); }
    w.u32((uint32_t)t.leafBrushes.size());
    for (uint32_t i : t.leafBrushes) w.u32(i);
    w.u32((uint32_t)t.brushes.size());
    for (const CollisionBrush& b : t.brushes) writeCollisionBrush(w, b);
    w.u32((uint32_t)cm.collMeshes.size());
    for (const CollisionMesh& c : cm.collMeshes) {
        w.u32(c.objectId); w.aabb(c.bounds);
        w.u32((uint32_t)c.tris.size());
        for (const Vec3& v : c.tris) w.vec3(v);
    }
    w.u32((uint32_t)cm.lights.size());
    for (const Light& l : cm.lights) {
        w.u8((uint8_t)l.type); w.vec3(l.pos); w.vec3(l.dir); w.vec3(l.color);
        w.f32(l.intensity); w.f32(l.radius); w.f32(l.innerDeg); w.f32(l.outerDeg); w.u32(l.entityId);
    }
}

// ------------------------------------------------------------ readers
void readMats(Reader& r, std::vector<Material>& out) {
    uint32_t c = r.count(36);
    out.resize(c);
    for (Material& m : out) {
        m.name = r.str(); m.texture = r.str(); m.tileSize = r.vec2(); m.color = r.vec3(); m.flags = r.u32();
    }
}

void readBrushes(Reader& r, std::vector<Brush>& out) {
    uint32_t c = r.count(20);
    out.resize(c);
    for (Brush& b : out) {
        b.id = r.u32(); b.owner = r.u32(); b.contents = r.u32(); b.flags = r.u32();
        uint32_t nf = r.count(44);
        b.faces.resize(nf);
        for (Face& f : b.faces) { f.plane = r.plane(); f.material = r.i32(); f.uv = r.uv(); f.flags = r.u32(); }
    }
}

void readMeshes(Reader& r, std::vector<MeshObject>& out) {
    uint32_t c = r.count(24);
    out.resize(c);
    for (MeshObject& m : out) {
        m.id = r.u32(); m.owner = r.u32(); m.flags = r.u32(); m.name = r.str();
        uint32_t nv = r.count(12);
        m.verts.resize(nv);
        for (Vec3& v : m.verts) v = r.vec3();
        uint32_t np = r.count(32);
        m.polys.resize(np);
        for (MeshPoly& p : m.polys) {
            p.material = r.i32(); p.smooth = r.u32(); p.proj = r.uv();
            uint32_t ni = r.count(4);
            p.idx.resize(ni);
            for (uint32_t& i : p.idx) i = r.u32();
            uint32_t nu = r.count(8);
            p.uv.resize(nu);
            for (Vec2& t : p.uv) t = r.vec2();
        }
    }
}

void readEntities(Reader& r, std::vector<Entity>& out) {
    uint32_t c = r.count(52);
    out.resize(c);
    for (Entity& e : out) {
        e.id = r.u32(); e.flags = r.u32(); e.classname = r.str();
        e.origin = r.vec3(); e.angles = r.vec3(); e.scale = r.vec3();
        uint32_t nk = r.count(8);
        e.props.resize(nk);
        for (KeyValue& kv : e.props) { kv.key = r.str(); kv.value = r.str(); }
    }
}

void readCollisionBrush(Reader& r, CollisionBrush& b) {
    b.contents = r.u32(); b.sourceId = r.u32(); b.bounds = r.aabb();
    uint32_t n = r.count(16);
    b.planes.resize(n);
    for (Plane& p : b.planes) p = r.plane();
}

void readCompiledChunk(Reader& r, CompiledMap& cm) {
    cm.ambient = r.vec3(); cm.unitsPerMeter = r.f32(); cm.baked = r.u8() != 0; cm.bounds = r.aabb();
    uint32_t nb = r.count(17);
    cm.batches.resize(nb);
    for (Batch& b : cm.batches) {
        b.material = r.i32(); b.source = (GeomSource)r.u8(); b.model = r.i32();
        uint32_t nv = r.count(36);
        b.vertices.resize(nv);
        for (Vertex& v : b.vertices) { v.pos = r.vec3(); v.normal = r.vec3(); v.uv = r.vec2(); v.color = r.u32(); b.bounds.add(v.pos); }
        uint32_t ni = r.count(4);
        b.indices.resize(ni);
        for (uint32_t& i : b.indices) { i = r.u32(); if (i >= nv) r.ok = false; }
    }
    uint32_t nm = r.count(40);
    cm.models.resize(nm);
    for (CompiledModel& m : cm.models) {
        m.entityId = r.u32(); m.bounds = r.aabb();
        uint32_t n = r.count(40);
        m.brushes.resize(n);
        for (CollisionBrush& b : m.brushes) readCollisionBrush(r, b);
    }
    Bsp& t = cm.bsp;
    t.root = r.i32(); t.depth = r.u32(); t.bounds = r.aabb();
    uint32_t np = r.count(16);
    t.planes.resize(np);
    for (Plane& p : t.planes) p = r.plane();
    uint32_t nn = r.count(12);
    t.nodes.resize(nn);
    for (BspNode& n : t.nodes) { n.plane = r.i32(); n.child[0] = r.i32(); n.child[1] = r.i32(); }
    uint32_t nl = r.count(8);
    t.leaves.resize(nl);
    for (BspLeaf& l : t.leaves) { l.firstBrush = r.u32(); l.brushCount = r.u32(); }
    uint32_t nlb = r.count(4);
    t.leafBrushes.resize(nlb);
    for (uint32_t& i : t.leafBrushes) i = r.u32();
    uint32_t nbr = r.count(40);
    t.brushes.resize(nbr);
    for (CollisionBrush& b : t.brushes) readCollisionBrush(r, b);
    uint32_t ncm = r.count(32);
    cm.collMeshes.resize(ncm);
    for (CollisionMesh& c : cm.collMeshes) {
        c.objectId = r.u32(); c.bounds = r.aabb();
        uint32_t n = r.count(12);
        c.tris.resize(n);
        for (Vec3& v : c.tris) v = r.vec3();
    }
    uint32_t nlt = r.count(53);
    cm.lights.resize(nlt);
    for (Light& l : cm.lights) {
        l.type = (LightType)r.u8(); l.pos = r.vec3(); l.dir = r.vec3(); l.color = r.vec3();
        l.intensity = r.f32(); l.radius = r.f32(); l.innerDeg = r.f32(); l.outerDeg = r.f32(); l.entityId = r.u32();
    }
    // validate the tree so a corrupt file cannot make queries read out of bounds
    if (!r.ok) return;
    auto childOk = [&](int32_t c) { return c >= 0 ? (size_t)c < t.nodes.size() : (size_t)(~c) < t.leaves.size(); };
    bool good = childOk(t.root) && !t.leaves.empty();
    for (const BspNode& n : t.nodes)
        good = good && n.plane >= 0 && (size_t)n.plane < t.planes.size() && childOk(n.child[0]) && childOk(n.child[1]);
    for (const BspLeaf& l : t.leaves) good = good && (size_t)l.firstBrush + l.brushCount <= t.leafBrushes.size();
    for (uint32_t i : t.leafBrushes) good = good && i < t.brushes.size();
    if (!good) r.ok = false;
    for (const Batch& b : cm.batches) cm.stats.triangles += (uint32_t)(b.indices.size() / 3), cm.stats.vertices += (uint32_t)b.vertices.size();
    cm.stats.bspNodes = nn; cm.stats.bspLeaves = nl; cm.stats.bspDepth = t.depth; cm.stats.lights = nlt;
}

constexpr char kMagic[4] = {'H', 'M', 'A', 'P'};

Status writeFile(const std::string& path, const std::vector<uint8_t>& data) {
    std::ofstream f(path, std::ios::binary | std::ios::trunc);
    if (!f) return Status::Fail("cannot write " + path);
    f.write((const char*)data.data(), (std::streamsize)data.size());
    return f ? Status::Ok() : Status::Fail("write failed: " + path);
}

bool readFile(const std::string& path, std::vector<uint8_t>& out) {
    std::ifstream f(path, std::ios::binary | std::ios::ate);
    if (!f) return false;
    std::streamsize sz = f.tellg();
    if (sz < 0) return false;
    f.seekg(0);
    out.resize((size_t)sz);
    if (sz > 0) f.read((char*)out.data(), sz);
    return (bool)f;
}

Status parse(const uint8_t* data, size_t size, Map& map, CompiledMap* compiled, bool* hasCompiled) {
    Reader hdr(data, size);
    char magic[4];
    for (char& c : magic) c = (char)hdr.u8();
    uint32_t version = hdr.u32();
    hdr.u32();  // flags (informational)
    if (!hdr.ok || std::memcmp(magic, kMagic, 4) != 0) return Status::Fail("not an .hmap file");
    if (version > kFormatVersion) return Status::Fail("file was written by a newer version of Hammer++");
    size_t pos = hdr.pos;
    *hasCompiled = false;
    CompiledMap cm;
    while (pos + 8 <= size) {
        std::string tag((const char*)data + pos, 4);
        Reader sz(data + pos + 4, 4);
        uint32_t len = sz.u32();
        pos += 8;
        if ((size_t)len > size - pos) return Status::Fail("truncated chunk '" + tag + "'");
        Reader r(data + pos, len);
        if (tag == "META") {
            map.name = r.str(); map.settings.ambient = r.vec3(); map.settings.gridSize = r.f32();
            map.settings.unitsPerMeter = r.f32(); map.nextId = r.u32();
        } else if (tag == "MATS") readMats(r, map.materials);
        else if (tag == "BRSH") readBrushes(r, map.brushes);
        else if (tag == "MESH") readMeshes(r, map.meshes);
        else if (tag == "ENTS") readEntities(r, map.entities);
        else if (tag == "CMPL") { readCompiledChunk(r, cm); *hasCompiled = r.ok; }
        if (!r.ok) return Status::Fail("corrupt chunk '" + tag + "'");
        pos += len;
    }
    // keep id allocation safe even when META is missing
    for (const Brush& b : map.brushes) map.nextId = std::max(map.nextId, b.id + 1);
    for (const MeshObject& m : map.meshes) map.nextId = std::max(map.nextId, m.id + 1);
    for (const Entity& e : map.entities) map.nextId = std::max(map.nextId, e.id + 1);
    if (*hasCompiled && compiled) {
        cm.materials = map.materials;
        cm.entities = map.entities;
        *compiled = std::move(cm);
    }
    return Status::Ok();
}

}  // namespace

std::vector<uint8_t> saveMapToMemory(const Map& map, const CompiledMap* compiled) {
    Writer out;
    out.buf.insert(out.buf.end(), kMagic, kMagic + 4);
    out.u32(kFormatVersion);
    out.u32(1u | (compiled ? 2u : 0u));
    Writer meta;
    meta.str(map.name); meta.vec3(map.settings.ambient); meta.f32(map.settings.gridSize);
    meta.f32(map.settings.unitsPerMeter); meta.u32(map.nextId);
    out.raw("META", meta);
    Writer w;
    writeMats(w, map.materials); out.raw("MATS", w);
    Writer b; writeBrushes(b, map.brushes); out.raw("BRSH", b);
    Writer m; writeMeshes(m, map.meshes); out.raw("MESH", m);
    Writer e; writeEntities(e, map.entities); out.raw("ENTS", e);
    if (compiled) { Writer c; writeCompiledChunk(c, *compiled); out.raw("CMPL", c); }
    return std::move(out.buf);
}

Status loadMapFromMemory(const uint8_t* data, size_t size, Map& out, CompiledMap* compiled) {
    Map tmp;
    bool has = false;
    Status s = parse(data, size, tmp, compiled, &has);
    if (!s) return s;
    out = std::move(tmp);
    return Status::Ok();
}

Status saveMap(const std::string& path, const Map& map, const CompiledMap* compiled) {
    return writeFile(path, saveMapToMemory(map, compiled));
}

Status loadMap(const std::string& path, Map& out, CompiledMap* compiled) {
    std::vector<uint8_t> data;
    if (!readFile(path, data)) return Status::Fail("cannot read " + path);
    return loadMapFromMemory(data.data(), data.size(), out, compiled);
}

Status saveCompiled(const std::string& path, const CompiledMap& cm) {
    Writer out;
    out.buf.insert(out.buf.end(), kMagic, kMagic + 4);
    out.u32(kFormatVersion);
    out.u32(2u);
    Writer w; writeMats(w, cm.materials); out.raw("MATS", w);
    Writer e; writeEntities(e, cm.entities); out.raw("ENTS", e);
    Writer c; writeCompiledChunk(c, cm); out.raw("CMPL", c);
    return writeFile(path, out.buf);
}

Status loadCompiled(const std::string& path, CompiledMap& out) {
    std::vector<uint8_t> data;
    if (!readFile(path, data)) return Status::Fail("cannot read " + path);
    Map map;
    CompiledMap cm;
    bool has = false;
    Status s = parse(data.data(), data.size(), map, &cm, &has);
    if (!s) return s;
    out = has ? std::move(cm) : compile(map);
    return Status::Ok();
}

}  // namespace hmap