#include "cli_wom_io.hpp"

#include "pipeline/wowee_model.hpp"
#include <glm/glm.hpp>
#include <nlohmann/json.hpp>

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>

namespace wowee {
namespace editor {
namespace cli {

namespace {

int handleExportObj(int& i, int argc, char** argv) {
    // Convert WOM (our open M2 replacement) to Wavefront OBJ — a
    // universally supported text format that opens directly in
    // Blender, MeshLab, ZBrush, Maya, and basically every other 3D
    // tool ever made. Makes the open-format ecosystem actually
    // useful for content authors who don't want to write a custom
    // WOM importer for their DCC of choice.
    std::string base = argv[++i];
    std::string outPath;
    if (i + 1 < argc && argv[i + 1][0] != '-') {
        outPath = argv[++i];
    }
    if (base.size() >= 4 && base.substr(base.size() - 4) == ".wom")
        base = base.substr(0, base.size() - 4);
    if (!wowee::pipeline::WoweeModelLoader::exists(base)) {
        std::fprintf(stderr, "WOM not found: %s.wom\n", base.c_str());
        return 1;
    }
    if (outPath.empty()) outPath = base + ".obj";
    auto wom = wowee::pipeline::WoweeModelLoader::load(base);
    if (!wom.isValid()) {
        std::fprintf(stderr, "WOM has no geometry to export: %s.wom\n", base.c_str());
        return 1;
    }
    std::ofstream obj(outPath);
    if (!obj) {
        std::fprintf(stderr, "Failed to open output file: %s\n", outPath.c_str());
        return 1;
    }
    // Header — preserves provenance so a designer reopening the OBJ
    // weeks later knows where it came from. The MTL line is a
    // courtesy: we don't currently emit a .mtl, but downstream
    // tools won't error without one either.
    obj << "# Wavefront OBJ generated by wowee_editor --export-obj\n";
    obj << "# Source: " << base << ".wom (v" << wom.version << ")\n";
    obj << "# Verts: " << wom.vertices.size()
        << " Tris: " << wom.indices.size() / 3
        << " Textures: " << wom.texturePaths.size() << "\n\n";
    obj << "o " << (wom.name.empty() ? "WoweeModel" : wom.name) << "\n";
    // Positions (v), texcoords (vt), normals (vn) — OBJ flips V so
    // that the same UVs that look right in our Vulkan renderer
    // also look right in Blender's bottom-left UV convention.
    for (const auto& v : wom.vertices) {
        obj << "v " << v.position.x << " " << v.position.y
            << " " << v.position.z << "\n";
    }
    for (const auto& v : wom.vertices) {
        obj << "vt " << v.texCoord.x << " " << (1.0f - v.texCoord.y) << "\n";
    }
    for (const auto& v : wom.vertices) {
        obj << "vn " << v.normal.x << " " << v.normal.y
            << " " << v.normal.z << "\n";
    }
    // Faces — split per-batch so each material/texture range becomes
    // its own group. Falls back to a single group when the WOM
    // wasn't authored with batches (WOM1/WOM2). OBJ indices are
    // 1-based, hence the +1.
    auto emitFaces = [&](const char* groupName,
                          uint32_t start, uint32_t count) {
        obj << "g " << groupName << "\n";
        for (uint32_t k = 0; k < count; k += 3) {
            uint32_t i0 = wom.indices[start + k] + 1;
            uint32_t i1 = wom.indices[start + k + 1] + 1;
            uint32_t i2 = wom.indices[start + k + 2] + 1;
            obj << "f "
                << i0 << "/" << i0 << "/" << i0 << " "
                << i1 << "/" << i1 << "/" << i1 << " "
                << i2 << "/" << i2 << "/" << i2 << "\n";
        }
    };
    if (wom.batches.empty()) {
        emitFaces("mesh", 0,
                  static_cast<uint32_t>(wom.indices.size()));
    } else {
        for (size_t b = 0; b < wom.batches.size(); ++b) {
            const auto& batch = wom.batches[b];
            std::string groupName = "batch_" + std::to_string(b);
            if (batch.textureIndex < wom.texturePaths.size()) {
                // Strip directory + extension for a readable group
                // name; full path is preserved in the file header
                // comment so nothing is lost.
                std::string tex = wom.texturePaths[batch.textureIndex];
                auto slash = tex.find_last_of("/\\");
                if (slash != std::string::npos) tex = tex.substr(slash + 1);
                auto dot = tex.find_last_of('.');
                if (dot != std::string::npos) tex = tex.substr(0, dot);
                if (!tex.empty()) groupName += "_" + tex;
            }
            emitFaces(groupName.c_str(), batch.indexStart, batch.indexCount);
        }
    }
    obj.close();
    std::printf("Exported %s.wom -> %s\n", base.c_str(), outPath.c_str());
    std::printf("  %zu verts, %zu tris, %zu groups\n",
                wom.vertices.size(), wom.indices.size() / 3,
                wom.batches.empty() ? size_t(1) : wom.batches.size());
    return 0;
}

int handleExportGlb(int& i, int argc, char** argv) {
    // glTF 2.0 binary (.glb) export — modern industry standard
    // that, unlike OBJ, supports skinning + animations + PBR
    // materials natively. v1 here writes positions/normals/UVs/
    // indices as a single mesh (or one primitive per WOM3 batch);
    // bones/anims are deliberately not yet emitted because glTF's
    // joint matrix layout differs from WOM's bone tree and needs
    // a careful re-mapping pass.
    //
    // Why this matters: glTF is what Sketchfab, Three.js, Babylon.js,
    // and Unity/Unreal-via-import all consume. Shipping WOM through
    // .glb makes our open binary format viewable in any modern
    // browser-based 3D viewer with zero conversion friction.
    std::string base = argv[++i];
    std::string outPath;
    if (i + 1 < argc && argv[i + 1][0] != '-') {
        outPath = argv[++i];
    }
    if (base.size() >= 4 && base.substr(base.size() - 4) == ".wom")
        base = base.substr(0, base.size() - 4);
    if (!wowee::pipeline::WoweeModelLoader::exists(base)) {
        std::fprintf(stderr, "WOM not found: %s.wom\n", base.c_str());
        return 1;
    }
    if (outPath.empty()) outPath = base + ".glb";
    auto wom = wowee::pipeline::WoweeModelLoader::load(base);
    if (!wom.isValid()) {
        std::fprintf(stderr, "WOM has no geometry: %s.wom\n", base.c_str());
        return 1;
    }
    // BIN chunk layout — sections ordered so each accessor's
    // byteOffset is naturally aligned for its component type:
    //   positions (vec3 float)  : 12 bytes/vert, offset 0
    //   normals   (vec3 float)  : 12 bytes/vert
    //   uvs       (vec2 float)  :  8 bytes/vert
    //   indices   (uint32)      :  4 bytes each
    // After 32 bytes per vertex, indices start at a 4-byte aligned
    // offset for free.
    const uint32_t vCount = static_cast<uint32_t>(wom.vertices.size());
    const uint32_t iCount = static_cast<uint32_t>(wom.indices.size());
    const uint32_t posOff = 0;
    const uint32_t nrmOff = posOff + vCount * 12;
    const uint32_t uvOff  = nrmOff + vCount * 12;
    const uint32_t idxOff = uvOff  + vCount * 8;
    const uint32_t binSize = idxOff + iCount * 4;
    std::vector<uint8_t> bin(binSize);
    // Pack positions
    for (uint32_t v = 0; v < vCount; ++v) {
        const auto& vert = wom.vertices[v];
        std::memcpy(&bin[posOff + v * 12 + 0], &vert.position.x, 4);
        std::memcpy(&bin[posOff + v * 12 + 4], &vert.position.y, 4);
        std::memcpy(&bin[posOff + v * 12 + 8], &vert.position.z, 4);
        std::memcpy(&bin[nrmOff + v * 12 + 0], &vert.normal.x, 4);
        std::memcpy(&bin[nrmOff + v * 12 + 4], &vert.normal.y, 4);
        std::memcpy(&bin[nrmOff + v * 12 + 8], &vert.normal.z, 4);
        std::memcpy(&bin[uvOff  + v * 8  + 0], &vert.texCoord.x, 4);
        std::memcpy(&bin[uvOff  + v * 8  + 4], &vert.texCoord.y, 4);
    }
    std::memcpy(&bin[idxOff], wom.indices.data(), iCount * 4);
    // Compute bounds for the position accessor's min/max — glTF
    // viewers rely on these for camera framing and culling.
    glm::vec3 bMin{1e30f}, bMax{-1e30f};
    for (const auto& v : wom.vertices) {
        bMin = glm::min(bMin, v.position);
        bMax = glm::max(bMax, v.position);
    }
    // Build the JSON structure. nlohmann::json keeps insertion
    // order in dump(), but glTF readers are key-based so order
    // doesn't matter functionally.
    nlohmann::json gj;
    gj["asset"] = {{"version", "2.0"},
                    {"generator", "wowee_editor --export-glb"}};
    gj["scene"] = 0;
    gj["scenes"] = nlohmann::json::array({nlohmann::json{{"nodes", {0}}}});
    gj["nodes"] = nlohmann::json::array({nlohmann::json{
        {"name", wom.name.empty() ? "WoweeModel" : wom.name},
        {"mesh", 0}
    }});
    gj["buffers"] = nlohmann::json::array({nlohmann::json{
        {"byteLength", binSize}
    }});
    // BufferViews: one per attribute + one per index range.
    // Per WOM3 batch we slice the index bufferView with separate
    // accessors so each batch becomes its own primitive.
    nlohmann::json bufferViews = nlohmann::json::array();
    // 0: positions, 1: normals, 2: uvs, 3: indices (whole range)
    bufferViews.push_back({{"buffer", 0}, {"byteOffset", posOff},
                            {"byteLength", vCount * 12},
                            {"target", 34962}}); // ARRAY_BUFFER
    bufferViews.push_back({{"buffer", 0}, {"byteOffset", nrmOff},
                            {"byteLength", vCount * 12},
                            {"target", 34962}});
    bufferViews.push_back({{"buffer", 0}, {"byteOffset", uvOff},
                            {"byteLength", vCount * 8},
                            {"target", 34962}});
    bufferViews.push_back({{"buffer", 0}, {"byteOffset", idxOff},
                            {"byteLength", iCount * 4},
                            {"target", 34963}}); // ELEMENT_ARRAY_BUFFER
    gj["bufferViews"] = bufferViews;
    // Accessors: 0=position, 1=normal, 2=uv, 3..N=indices (one
    // per primitive, sliced from bufferView 3).
    nlohmann::json accessors = nlohmann::json::array();
    accessors.push_back({
        {"bufferView", 0}, {"componentType", 5126}, // FLOAT
        {"count", vCount}, {"type", "VEC3"},
        {"min", {bMin.x, bMin.y, bMin.z}},
        {"max", {bMax.x, bMax.y, bMax.z}}
    });
    accessors.push_back({
        {"bufferView", 1}, {"componentType", 5126},
        {"count", vCount}, {"type", "VEC3"}
    });
    accessors.push_back({
        {"bufferView", 2}, {"componentType", 5126},
        {"count", vCount}, {"type", "VEC2"}
    });
    // Build primitives — one per WOM3 batch, or one over the
    // whole index range if no batches.
    nlohmann::json primitives = nlohmann::json::array();
    auto addPrimitive = [&](uint32_t idxStart, uint32_t idxCount) {
        uint32_t accessorIdx = static_cast<uint32_t>(accessors.size());
        accessors.push_back({
            {"bufferView", 3},
            {"byteOffset", idxStart * 4},
            {"componentType", 5125}, // UNSIGNED_INT
            {"count", idxCount},
            {"type", "SCALAR"}
        });
        primitives.push_back({
            {"attributes", {{"POSITION", 0}, {"NORMAL", 1}, {"TEXCOORD_0", 2}}},
            {"indices", accessorIdx},
            {"mode", 4} // TRIANGLES
        });
    };
    if (wom.batches.empty()) {
        addPrimitive(0, iCount);
    } else {
        for (const auto& b : wom.batches) {
            addPrimitive(b.indexStart, b.indexCount);
        }
    }
    gj["accessors"] = accessors;
    gj["meshes"] = nlohmann::json::array({nlohmann::json{
        {"primitives", primitives}
    }});
    // Serialize JSON to bytes; pad to 4-byte boundary with spaces
    // (glTF spec requires JSON chunk padded with 0x20).
    std::string jsonStr = gj.dump();
    while (jsonStr.size() % 4 != 0) jsonStr += ' ';
    // BIN chunk pads to 4-byte boundary with zeros (already
    // satisfied since binSize = idxOff + iCount*4 and idxOff is
    // 4-byte aligned).
    uint32_t jsonLen = static_cast<uint32_t>(jsonStr.size());
    uint32_t binLen = binSize;
    uint32_t totalLen = 12 + 8 + jsonLen + 8 + binLen;
    std::ofstream out(outPath, std::ios::binary);
    if (!out) {
        std::fprintf(stderr, "Failed to open output: %s\n", outPath.c_str());
        return 1;
    }
    // Header: magic, version, total length (all little-endian uint32)
    uint32_t magic = 0x46546C67;  // 'glTF'
    uint32_t version = 2;
    out.write(reinterpret_cast<const char*>(&magic), 4);
    out.write(reinterpret_cast<const char*>(&version), 4);
    out.write(reinterpret_cast<const char*>(&totalLen), 4);
    // JSON chunk header + payload
    uint32_t jsonChunkType = 0x4E4F534A;  // 'JSON'
    out.write(reinterpret_cast<const char*>(&jsonLen), 4);
    out.write(reinterpret_cast<const char*>(&jsonChunkType), 4);
    out.write(jsonStr.data(), jsonLen);
    // BIN chunk header + payload
    uint32_t binChunkType = 0x004E4942;  // 'BIN\0'
    out.write(reinterpret_cast<const char*>(&binLen), 4);
    out.write(reinterpret_cast<const char*>(&binChunkType), 4);
    out.write(reinterpret_cast<const char*>(bin.data()), binLen);
    out.close();
    std::printf("Exported %s.wom -> %s\n", base.c_str(), outPath.c_str());
    std::printf("  %u verts, %u tris, %zu primitive(s), %u-byte binary chunk\n",
                vCount, iCount / 3, primitives.size(), binLen);
    return 0;
}

int handleExportStl(int& i, int argc, char** argv) {
    // ASCII STL export — single most universal 3D-printer format.
    // Cura, PrusaSlicer, Bambu Studio, Slic3r, OctoPrint, MakerBot
    // — every slicer made in the last 25 years opens STL natively.
    // Lets WOM models drive physical prints with no conversion
    // friction beyond this one command.
    std::string base = argv[++i];
    std::string outPath;
    if (i + 1 < argc && argv[i + 1][0] != '-') outPath = argv[++i];
    if (base.size() >= 4 && base.substr(base.size() - 4) == ".wom")
        base = base.substr(0, base.size() - 4);
    if (!wowee::pipeline::WoweeModelLoader::exists(base)) {
        std::fprintf(stderr, "WOM not found: %s.wom\n", base.c_str());
        return 1;
    }
    if (outPath.empty()) outPath = base + ".stl";
    auto wom = wowee::pipeline::WoweeModelLoader::load(base);
    if (!wom.isValid()) {
        std::fprintf(stderr, "WOM has no geometry: %s.wom\n", base.c_str());
        return 1;
    }
    std::ofstream out(outPath);
    if (!out) {
        std::fprintf(stderr, "Failed to open output: %s\n", outPath.c_str());
        return 1;
    }
    // STL solid name must be alphanumeric + underscores per loose
    // convention; sanitize whatever the WOM name contains. Empty
    // -> 'wowee_model'.
    std::string solidName = wom.name.empty() ? "wowee_model" : wom.name;
    for (auto& c : solidName) {
        if (!((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
              (c >= '0' && c <= '9') || c == '_')) c = '_';
    }
    out << "solid " << solidName << "\n";
    // Per-triangle facet — STL has no shared vertex pool, every
    // triangle stands alone. Compute face normal from cross product
    // (STL spec requires unit-length face normal; viewers fall
    // back to per-vertex if zero, but most slicers want the real
    // value for orientation hints).
    uint32_t triCount = 0;
    for (size_t k = 0; k + 2 < wom.indices.size(); k += 3) {
        uint32_t i0 = wom.indices[k];
        uint32_t i1 = wom.indices[k + 1];
        uint32_t i2 = wom.indices[k + 2];
        if (i0 >= wom.vertices.size() || i1 >= wom.vertices.size() ||
            i2 >= wom.vertices.size()) continue;
        const auto& v0 = wom.vertices[i0].position;
        const auto& v1 = wom.vertices[i1].position;
        const auto& v2 = wom.vertices[i2].position;
        glm::vec3 e1 = v1 - v0;
        glm::vec3 e2 = v2 - v0;
        glm::vec3 n = glm::cross(e1, e2);
        float len = glm::length(n);
        if (len > 1e-12f) n /= len;
        else n = {0, 0, 1};  // degenerate — STL spec allows any unit normal
        out << "  facet normal " << n.x << " " << n.y << " " << n.z << "\n"
            << "    outer loop\n"
            << "      vertex " << v0.x << " " << v0.y << " " << v0.z << "\n"
            << "      vertex " << v1.x << " " << v1.y << " " << v1.z << "\n"
            << "      vertex " << v2.x << " " << v2.y << " " << v2.z << "\n"
            << "    endloop\n"
            << "  endfacet\n";
        triCount++;
    }
    out << "endsolid " << solidName << "\n";
    out.close();
    std::printf("Exported %s.wom -> %s\n", base.c_str(), outPath.c_str());
    std::printf("  solid '%s', %u facets\n",
                solidName.c_str(), triCount);
    return 0;
}

int handleImportStl(int& i, int argc, char** argv) {
    // ASCII STL -> WOM. Closes the STL round trip so designers can
    // edit prints in TinkerCAD/Meshmixer/SolidWorks and bring them
    // back to the engine. Dedupes vertices on (pos, normal) so the
    // resulting WOM vertex buffer stays compact.
    std::string stlPath = argv[++i];
    std::string womBase;
    if (i + 1 < argc && argv[i + 1][0] != '-') womBase = argv[++i];
    if (!std::filesystem::exists(stlPath)) {
        std::fprintf(stderr, "STL not found: %s\n", stlPath.c_str());
        return 1;
    }
    if (womBase.empty()) {
        womBase = stlPath;
        if (womBase.size() >= 4 &&
            womBase.substr(womBase.size() - 4) == ".stl") {
            womBase = womBase.substr(0, womBase.size() - 4);
        }
    }
    std::ifstream in(stlPath);
    if (!in) {
        std::fprintf(stderr, "Failed to open STL: %s\n", stlPath.c_str());
        return 1;
    }
    wowee::pipeline::WoweeModel wom;
    wom.version = 1;
    // Dedupe key: 6 floats (pos + normal) packed as a string. Loose
    // matching, but exact for round-trips since we write the same
    // floats back. Real-world STLs from CAD tools rarely benefit
    // from looser tolerance — they already share verts at the
    // exporter level.
    std::unordered_map<std::string, uint32_t> dedupe;
    auto interVert = [&](const glm::vec3& pos, const glm::vec3& nrm) {
        char key[128];
        std::snprintf(key, sizeof(key), "%.6f|%.6f|%.6f|%.6f|%.6f|%.6f",
                      pos.x, pos.y, pos.z, nrm.x, nrm.y, nrm.z);
        auto it = dedupe.find(key);
        if (it != dedupe.end()) return it->second;
        wowee::pipeline::WoweeModel::Vertex v;
        v.position = pos;
        v.normal = nrm;
        v.texCoord = {0, 0};
        uint32_t idx = static_cast<uint32_t>(wom.vertices.size());
        wom.vertices.push_back(v);
        dedupe[key] = idx;
        return idx;
    };
    std::string line;
    std::string solidName;
    // Per-facet state: parsed normal + accumulating vertex queue.
    glm::vec3 currentNormal{0, 0, 1};
    std::vector<glm::vec3> facetVerts;
    int facetCount = 0;
    while (std::getline(in, line)) {
        while (!line.empty() && (line.back() == '\r' || line.back() == ' '))
            line.pop_back();
        std::istringstream ss(line);
        std::string tok;
        ss >> tok;
        if (tok == "solid" && solidName.empty()) {
            ss >> solidName;
        } else if (tok == "facet") {
            std::string normalKw;
            ss >> normalKw;
            if (normalKw == "normal") {
                ss >> currentNormal.x >> currentNormal.y >> currentNormal.z;
            }
            facetVerts.clear();
        } else if (tok == "vertex") {
            glm::vec3 v;
            ss >> v.x >> v.y >> v.z;
            facetVerts.push_back(v);
        } else if (tok == "endfacet") {
            if (facetVerts.size() == 3) {
                // Use the facet normal for all 3 verts since STL
                // doesn't carry per-vertex normals. Glue-points to
                // adjacent facets will get distinct verts (which is
                // correct for faceted-shading STL geometry).
                for (const auto& v : facetVerts) {
                    wom.indices.push_back(interVert(v, currentNormal));
                }
                facetCount++;
            }
            facetVerts.clear();
        }
        // 'outer loop', 'endloop', 'endsolid' ignored — we infer
        // from the vertex count per facet.
    }
    if (wom.vertices.empty() || wom.indices.empty()) {
        std::fprintf(stderr,
            "import-stl: no geometry parsed from %s\n", stlPath.c_str());
        return 1;
    }
    wom.name = solidName.empty()
        ? std::filesystem::path(stlPath).stem().string()
        : solidName;
    // Compute bounds — renderer culls by these so wrong values
    // make models disappear at distance.
    wom.boundMin = wom.vertices[0].position;
    wom.boundMax = wom.boundMin;
    for (const auto& v : wom.vertices) {
        wom.boundMin = glm::min(wom.boundMin, v.position);
        wom.boundMax = glm::max(wom.boundMax, v.position);
    }
    glm::vec3 center = (wom.boundMin + wom.boundMax) * 0.5f;
    float r2 = 0;
    for (const auto& v : wom.vertices) {
        glm::vec3 d = v.position - center;
        r2 = std::max(r2, glm::dot(d, d));
    }
    wom.boundRadius = std::sqrt(r2);
    if (!wowee::pipeline::WoweeModelLoader::save(wom, womBase)) {
        std::fprintf(stderr, "import-stl: failed to write %s.wom\n",
                     womBase.c_str());
        return 1;
    }
    std::printf("Imported %s -> %s.wom\n", stlPath.c_str(), womBase.c_str());
    std::printf("  %d facets, %zu verts (deduped), bounds [%.2f, %.2f, %.2f] - [%.2f, %.2f, %.2f]\n",
                facetCount, wom.vertices.size(),
                wom.boundMin.x, wom.boundMin.y, wom.boundMin.z,
                wom.boundMax.x, wom.boundMax.y, wom.boundMax.z);
    return 0;
}

}  // namespace

bool handleWomIo(int& i, int argc, char** argv, int& outRc) {
    if (std::strcmp(argv[i], "--export-obj") == 0 && i + 1 < argc) {
        outRc = handleExportObj(i, argc, argv); return true;
    }
    if (std::strcmp(argv[i], "--export-glb") == 0 && i + 1 < argc) {
        outRc = handleExportGlb(i, argc, argv); return true;
    }
    if (std::strcmp(argv[i], "--export-stl") == 0 && i + 1 < argc) {
        outRc = handleExportStl(i, argc, argv); return true;
    }
    if (std::strcmp(argv[i], "--import-stl") == 0 && i + 1 < argc) {
        outRc = handleImportStl(i, argc, argv); return true;
    }
    return false;
}

} // namespace cli
} // namespace editor
} // namespace wowee