#include "visualisation/handler_import_3mf.h"
#include "util/verbose.h"
#include <string>
#include <vector>
#include <tinyxml2.h>
#include <minizip/unzip.h>
#include <minizip/zip.h>
#include <minizip/ioapi.h>
#include "shape/shape.h"
#include "parser/evaluate_expression.h"
#include "parser/value.h"
#include "manifold/manifold.h"
#include <fstream>
#include <unordered_map>
#include <memory>
#include "util/obj_utils.h"

using namespace tinyxml2;

// Helper: Extracts XML from any file in the zip
std::string extract_xml_from_zip(const std::string& zip_path, const std::string& inner_path) {
    unzFile zipfile = unzOpen64(zip_path.c_str());
    if (!zipfile) return "";
    if (unzLocateFile(zipfile, inner_path.c_str(), 0) != UNZ_OK) {
        unzClose(zipfile);
        return "";
    }
    if (unzOpenCurrentFile(zipfile) != UNZ_OK) {
        unzClose(zipfile);
        return "";
    }
    unz_file_info file_info;
    if (unzGetCurrentFileInfo(zipfile, &file_info, nullptr, 0, nullptr, 0, nullptr, 0) != UNZ_OK) {
        unzCloseCurrentFile(zipfile);
        unzClose(zipfile);
        return "";
    }
    std::string xml_content;
    xml_content.resize(file_info.uncompressed_size);
    int bytes_read = unzReadCurrentFile(zipfile, &xml_content[0], file_info.uncompressed_size);
    unzCloseCurrentFile(zipfile);
    unzClose(zipfile);
    if (bytes_read <= 0) return "";
    xml_content.resize(bytes_read);
    return xml_content;
}

// Helper: Enumerate all .model files in the zip
std::vector<std::string> list_model_files_in_zip(const std::string& zip_path) {
    std::vector<std::string> model_files;
    unzFile zipfile = unzOpen64(zip_path.c_str());
    if (!zipfile) return model_files;
    if (unzGoToFirstFile(zipfile) != UNZ_OK) {
        unzClose(zipfile);
        return model_files;
    }
    do {
        char filename[512];
        unz_file_info file_info;
        if (unzGetCurrentFileInfo(zipfile, &file_info, filename, sizeof(filename), nullptr, 0, nullptr, 0) == UNZ_OK) {
            std::string fname(filename);
            if (fname.size() > 6 && fname.substr(fname.size() - 6) == ".model") {
                model_files.push_back(fname);
            }
        }
    } while (unzGoToNextFile(zipfile) == UNZ_OK);
    unzClose(zipfile);
    return model_files;
}

// Helper: Extracts /3D/3dmodel.model from 3mf zip
std::string extract_3dmodel_xml(const std::string& path_to_3mf) {
    return extract_xml_from_zip(path_to_3mf, "3D/3dmodel.model");
}

// Helper: Parses all mesh objects from all .model XMLs and returns a vector of Shapes
std::vector<Shape> parse_3mf_meshes(const std::string& xml_content, const std::string& filename) {
    verbose_log("import_3mf", "Parsing 3MF XML");
    XMLDocument doc;
    std::vector<Shape> shapes;
    if (doc.Parse(xml_content.c_str()) != XML_SUCCESS) {
        verbose_log("import_3mf", "Failed to parse XML");
        return shapes;
    }

    XMLElement* model = doc.FirstChildElement("model");
    if (!model) model = doc.FirstChildElement("m:model"); // handle namespace
    if (!model) {
        verbose_log("import_3mf", "No <model> element found");
        return shapes;
    }

    // --- Collect object IDs referenced in <build> ---
    std::vector<int> build_ids;
    XMLElement* build = model->FirstChildElement("build");
    if (build) {
        for (XMLElement* item = build->FirstChildElement("item"); item; item = item->NextSiblingElement("item")) {
            const char* oid = item->Attribute("objectid");
            if (oid) {
                int id = std::atoi(oid);
                build_ids.push_back(id);
                verbose_log("import_3mf", std::string("build references objectid=") + oid);
            }
        }
    }

    XMLElement* resources = model->FirstChildElement("resources");
    if (!resources) {
        verbose_log("import_3mf", "No <resources> element found");
        return shapes;
    }

    // --- Build a map of all objects by ID from ALL .model files in the zip ---
    std::unordered_map<int, XMLElement*> object_map;
    std::vector<std::unique_ptr<XMLDocument>> extra_docs;

    // 1. Add objects from main 3dmodel.model
    for (XMLElement* object = resources->FirstChildElement("object"); object; object = object->NextSiblingElement("object")) {
        const char* id_attr = object->Attribute("id");
        if (id_attr) {
            int id = std::atoi(id_attr);
            object_map[id] = object;
            const char* type_attr = object->Attribute("type");
            verbose_log("import_3mf", std::string("Found object id=") + id_attr +
                " type=" + (type_attr ? type_attr : "default"));
        }
    }

    // 2. Add objects from all /3D/Objects/*.model files
    std::vector<std::string> model_files = list_model_files_in_zip(filename);
    for (const std::string& fname : model_files) {
        // Skip the main 3dmodel.model (already loaded)
        if (fname == "3D/3dmodel.model") continue;
        std::string xml = extract_xml_from_zip(filename, fname);
        if (xml.empty()) continue;
        extra_docs.emplace_back(new XMLDocument());
        XMLDocument* doc_ptr = extra_docs.back().get();
        if (doc_ptr->Parse(xml.c_str()) != XML_SUCCESS) continue;
        XMLElement* model2 = doc_ptr->FirstChildElement("model");
        if (!model2) model2 = doc_ptr->FirstChildElement("m:model");
        if (!model2) continue;
        XMLElement* resources2 = model2->FirstChildElement("resources");
        if (!resources2) continue;
        for (XMLElement* object = resources2->FirstChildElement("object"); object; object = object->NextSiblingElement("object")) {
            const char* id_attr = object->Attribute("id");
            if (id_attr) {
                int id = std::atoi(id_attr);
                object_map[id] = object;
                const char* type_attr = object->Attribute("type");
                verbose_log("import_3mf", std::string("Found object id=") + id_attr +
                    " type=" + (type_attr ? type_attr : "default") + " (from " + fname + ")");
            }
        }
    }

    // --- Recursive mesh collector ---
    std::unordered_set<int> visited;
    std::vector<XMLElement*> mesh_objects;

    std::function<void(int)> collect_meshes = [&](int obj_id) {
        if (visited.count(obj_id)) return;
        visited.insert(obj_id);
        auto it = object_map.find(obj_id);
        if (it == object_map.end()) return;
        XMLElement* object = it->second;
        XMLElement* mesh = object->FirstChildElement("mesh");
        if (mesh) {
            mesh_objects.push_back(object);
        } else {
            XMLElement* components = object->FirstChildElement("components");
            if (components) {
                for (XMLElement* comp = components->FirstChildElement("component"); comp; comp = comp->NextSiblingElement("component")) {
                    const char* child_id_attr = comp->Attribute("objectid");
                    if (child_id_attr) {
                        int child_id = std::atoi(child_id_attr);
                        verbose_log("import_3mf", std::string("Recursing into component objectid=") + child_id_attr);
                        collect_meshes(child_id);
                    }
                }
            }
        }
    };

    // --- Start from all build objectids ---
    for (int obj_id : build_ids) {
        collect_meshes(obj_id);
    }

    int mesh_count = 0;
    for (XMLElement* object : mesh_objects) {
        XMLElement* mesh = object->FirstChildElement("mesh");
        if (!mesh) continue;

        // Parse vertices
        std::vector<std::array<float, 3>> vertices;
        XMLElement* verticesElem = mesh->FirstChildElement("vertices");
        if (!verticesElem) continue;
        for (XMLElement* v = verticesElem->FirstChildElement("vertex"); v; v = v->NextSiblingElement("vertex")) {
            float x = v->FloatAttribute("x");
            float y = v->FloatAttribute("y");
            float z = v->FloatAttribute("z");
            vertices.push_back({x, y, z});
        }

        // Parse triangles
        std::vector<std::array<int, 3>> triangles;
        XMLElement* trianglesElem = mesh->FirstChildElement("triangles");
        if (!trianglesElem) continue;
        for (XMLElement* t = trianglesElem->FirstChildElement("triangle"); t; t = t->NextSiblingElement("triangle")) {
            int v1 = t->IntAttribute("v1");
            int v2 = t->IntAttribute("v2");
            int v3 = t->IntAttribute("v3");
            triangles.push_back({v1, v2, v3});
        }

        verbose_log("import_3mf", "Loaded mesh " + std::to_string(mesh_count+1) +
            ": " + std::to_string(vertices.size()) + " vertices, " +
            std::to_string(triangles.size()) + " triangles");

        // Convert to internal mesh format using MeshGL and Manifold
        manifold::MeshGL meshgl;
        meshgl.numProp = 3;
        for (const auto& v : vertices) {
            meshgl.vertProperties.push_back(v[0]);
            meshgl.vertProperties.push_back(v[1]);
            meshgl.vertProperties.push_back(v[2]);
        }
        for (const auto& t : triangles) {
            meshgl.triVerts.push_back(static_cast<uint32_t>(t[0]));
            meshgl.triVerts.push_back(static_cast<uint32_t>(t[1]));
            meshgl.triVerts.push_back(static_cast<uint32_t>(t[2]));
        }
        manifold::Manifold man(meshgl);
        Shape s;
        s.mesh = std::move(man);
        s.origin_expr = "import_3mf_" + std::to_string(mesh_count+1);
        shapes.push_back(s);
        ++mesh_count;
    }
    verbose_log("import_3mf", "Total meshes loaded: " + std::to_string(mesh_count));
    return shapes;
}

Value handle_import_3mf(const PrimitiveCall& call) {
    verbose_log("import_3mf", "start");
    std::string filename;
    if (!call.positional.empty()) {
        filename = local_strip_quotes(local_trim(call.positional[0]));
    }
    verbose_log("import_3mf", "Importing file: " + filename);

    std::string xml = extract_3dmodel_xml(filename);
    if (xml.empty()) {
        verbose_log("import_3mf", "Failed to extract 3dmodel.model from 3mf");
        throw std::runtime_error("3MF extraction failed");
    }
    std::vector<Shape> meshes = parse_3mf_meshes(xml, filename);
    if (meshes.empty()) {
        verbose_log("import_3mf", "No meshes found in 3MF file");
        throw std::runtime_error("No meshes found in 3MF file");
    }

    // Print mesh names and counts
    verbose_log("import_3mf", "Imported " + std::to_string(meshes.size()) + " mesh(es):");
    for (size_t i = 0; i < meshes.size(); ++i) {
        std::string name = "obj3mf_" + std::to_string(i+1);
        verbose_log("import_3mf", "  [" + std::to_string(i) + "] " + name);
    }

    // Return as array of shapes (Value::Array)
    Value out;
    out.type = Value::Type::Array;
    for (auto& s : meshes) {
        Value v;
        v.type = Value::Type::Shape;
        v.shape = new Shape(std::move(s));
        out.array.push_back(v);
    }
    return out;
}

REGISTER_HELP(import_3mf,
    "Imports all meshes from a 3MF file and returns them as an array of shapes (no union).",
    "filename : string - Path to .3mf file\n",
    "// Import all meshes from a 3MF file as separate shapes\n"
    "parts = import_3mf(\"file.3mf\")\n"
    "show_viewer(parts) // shows all meshes\n"
    "save_obj(parts[0], \"obj3mf_1.obj\") // save first mesh\n"
    "save_obj(parts[1], \"obj3mf_2.obj\") // save second mesh\n",
    "1) Extracts /3D/3dmodel.model from the 3MF archive\n"
    "2) Parses all <object> meshes in the XML and all /3D/Objects/*.model files\n"
    "3) Converts each mesh to internal format (MeshGL → Manifold)\n"
    "4) Returns all meshes as an array of shapes (no union)\n"
    "5) Prints mesh names and count to the console",
    "All objects in /3D/3dmodel.model and /3D/Objects/*.model are included.\n"
    "Use show_viewer(parts) to view all meshes at once.\n"
    "Use save_obj(parts[i], \"obj3mf_N.obj\") to export each mesh individually.\n"
    "Mesh names: obj3mf_1, obj3mf_2, ...\n"
    "Verbose logging: verbose_set_module(import_3mf, true) for detailed import steps."
);