"""Streaming 3MF reader shared by convert_3mf.py, fingerprint.py and threemf_info.py.

trimesh's 3MF loader re-reads a whole production-extension part file (every
mesh in it) for each component that points into that file. Bambu Studio
projects are hundreds of components pointing into a few large part files, so
the geometry multiplies until the process runs out of memory.

This reader parses each part file once, element by element, resolves each
component to the one object it names, and returns every object's mesh once
plus the transforms where it is placed — so parts reused across plates share
their geometry. It also reads Bambu Studio / OrcaSlicer project metadata
(plates, filaments, part types) when the file has it.
"""
import json
import re
import zipfile
from array import array
from collections import namedtuple

import numpy as np
from lxml import etree

UNIT_MM = {"micron": 0.001, "millimeter": 1.0, "centimeter": 10.0, "inch": 25.4, "foot": 304.8, "meter": 1000.0}
PARSER = dict(huge_tree=True, resolve_entities=False, no_network=True)
MAX_DEPTH = 32
# Slicer part types that shape or steer printing but aren't printed themselves.
HIDDEN_SUBTYPES = {"negative_part", "modifier_part", "support_blocker", "support_enforcer"}

# One placed copy of an object's mesh. `item` is the object id of the build item
# it belongs to and `instance` which occurrence of that id it is; `part` is the
# component directly under that object (how slicer metadata names parts).
Placement = namedtuple("Placement", "key transform item instance part")


def _local(name):
    return name.rpartition("}")[2]


def _path_attr(el):
    # The production extension's p:path, whatever prefix the file bound it to.
    for key, value in el.attrib.items():
        if _local(key) == "path":
            return value
    return None


def _int(value):
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _matrix(text):
    m = np.eye(4)
    if text:
        v = [float(x) for x in text.split()]
        if len(v) == 12:
            # 3MF lists "m00 m01 m02 m10 ... m32" for row vectors; store column form.
            m[:3, :3] = np.array(v[:9]).reshape(3, 3).T
            m[:3, 3] = v[9:12]
    return m


def _parse_part(fh):
    """Returns (objects, build items, unit) for one model part.

    objects: {id: (vertices (n,3) float64, faces (m,3) int32, [(path|None, objectid, 4x4)])}
    """
    objects, items, unit = {}, [], None
    oid = verts = tris = comps = None
    for event, el in etree.iterparse(fh, events=("start", "end"), **PARSER):
        name = _local(el.tag)
        if event == "start":
            if name == "model":
                unit = el.get("unit")
            elif name == "object":
                oid, verts, tris, comps = el.get("id"), array("d"), array("i"), []
            continue

        if name == "vertex":
            verts.extend((float(el.get("x")), float(el.get("y")), float(el.get("z"))))
        elif name == "triangle":
            tris.extend((int(el.get("v1")), int(el.get("v2")), int(el.get("v3"))))
        elif name == "component":
            comps.append((_path_attr(el), el.get("objectid"), _matrix(el.get("transform"))))
        elif name == "item":
            items.append((el.get("objectid"), _matrix(el.get("transform"))))
        elif name == "object":
            v = np.frombuffer(verts, dtype=np.float64).reshape(-1, 3)
            f = np.frombuffer(tris, dtype=np.int32).reshape(-1, 3)
            if len(f) and (f.min() < 0 or f.max() >= len(v)):
                f = f[:0]  # corrupt indices: drop the mesh rather than crash later
            objects[oid] = (v, f, comps)

        # Free what's been read so memory stays flat on huge parts.
        el.clear()
        parent = el.getparent()
        while parent is not None and el.getprevious() is not None:
            del parent[0]
    return objects, items, unit


def _root_part(z):
    try:
        rels = etree.fromstring(z.read("_rels/.rels"), etree.XMLParser(resolve_entities=False, no_network=True))
        for rel in rels:
            if rel.get("Type", "").endswith("/3dmodel"):
                return rel.get("Target", "").lstrip("/")
    except (KeyError, etree.XMLSyntaxError):
        pass
    return "3D/3dmodel.model"


def load(path):
    """Returns (meshes, placements).

    meshes: {key: (vertices, faces)} — each object once, in file units.
    placements: [Placement] — one per placed instance; transforms map to millimetres.
    """
    with zipfile.ZipFile(path) as z:
        names = set(z.namelist())
        parts = {}

        def part(name):
            name = name.lstrip("/")
            if name not in parts:
                if name not in names:
                    raise ValueError(f"3MF references a missing part: {name}")
                with z.open(name) as fh:
                    parts[name] = _parse_part(fh)
            return parts[name]

        root = _root_part(z)
        _, items, unit = part(root)
        scale = np.diag([UNIT_MM.get(unit or "millimeter", 1.0)] * 3 + [1.0])
        meshes, placements = {}, []

        def walk(name, oid, transform, depth, item, instance, part_id):
            if depth > MAX_DEPTH:
                raise ValueError("3MF components nest too deeply")
            obj = part(name)[0].get(oid)
            if obj is None:
                return  # undefined reference; slicers skip these too
            v, f, comps = obj
            if len(f):
                key = f"{name}#{oid}"
                meshes.setdefault(key, (v, f))
                placements.append(Placement(key, transform, item, instance, part_id))
            for cpath, coid, m in comps:
                walk(cpath.lstrip("/") if cpath else name, coid, transform @ m, depth + 1,
                     item, instance, part_id if depth else coid)

        seen = {}
        for oid, m in items:
            instance = seen.get(oid, 0)
            seen[oid] = instance + 1
            walk(root, oid, scale @ m, 0, oid, instance, None)
    return meshes, placements


def _application(z, names):
    """The app that wrote the file, from the root model's <metadata name="Application">.
    Metadata comes before the geometry, so reading stops at <resources>."""
    root = _root_part(z)
    if root not in names:
        return None
    try:
        with z.open(root) as fh:
            for event, el in etree.iterparse(fh, events=("start", "end"), **PARSER):
                name = _local(el.tag)
                if event == "start":
                    if name in ("resources", "build"):
                        break
                elif name == "metadata" and el.get("name") == "Application":
                    return (el.text or "").strip() or None
    except etree.XMLSyntaxError:
        pass
    return None


def read_project(path):
    """Slicer project metadata (Bambu Studio / OrcaSlicer), empty when absent:

    {application: "BambuStudio-02.01.01.52" | None,
     plates: [{index, name, objects: [(object_id, instance)], thumbnail, thumbnail_small}],
     filaments: [{colour, type}],
     objects: {object_id: {name, extruder, parts: {part_id: {extruder, subtype}}}}}
    """
    project = {"application": None, "plates": [], "filaments": [], "objects": {}}
    with zipfile.ZipFile(path) as z:
        names = set(z.namelist())
        project["application"] = _application(z, names)

        if "Metadata/project_settings.config" in names:
            try:
                settings = json.loads(z.read("Metadata/project_settings.config"))
                colours = settings.get("filament_colour") or []
                types = settings.get("filament_type") or []
                project["filaments"] = [
                    {"colour": c, "type": types[i] if i < len(types) else None} for i, c in enumerate(colours)]
            except (ValueError, UnicodeDecodeError, AttributeError):
                pass

        if "Metadata/model_settings.config" in names:
            try:
                root = etree.fromstring(z.read("Metadata/model_settings.config"), etree.XMLParser(**PARSER))
            except etree.XMLSyntaxError:
                return project

            def meta(el):
                return {m.get("key"): m.get("value") for m in el.findall("metadata") if m.get("key")}

            for obj in root.findall("object"):
                info = meta(obj)
                project["objects"][obj.get("id")] = {
                    "name": info.get("name"),
                    "extruder": _int(info.get("extruder")),
                    "parts": {
                        p.get("id"): {"extruder": _int(meta(p).get("extruder")), "subtype": p.get("subtype") or "normal_part"}
                        for p in obj.findall("part")
                    },
                }
            for plate in root.findall("plate"):
                info = meta(plate)
                index = _int(info.get("plater_id")) or len(project["plates"]) + 1
                large = info.get("thumbnail_file") or f"Metadata/plate_{index}.png"
                small = f"Metadata/plate_{index}_small.png"
                project["plates"].append({
                    "index": index,
                    "name": info.get("plater_name") or "",
                    "objects": [(meta(i).get("object_id"), _int(meta(i).get("instance_id")) or 0)
                                for i in plate.findall("model_instance")],
                    "thumbnail": large if large in names else None,
                    "thumbnail_small": small if small in names else None,
                })
    return project


def printable(project, placements):
    """Drops negative, modifier and support parts (known from slicer metadata)."""
    def visible(p):
        obj = project["objects"].get(p.item)
        part = obj and obj["parts"].get(p.part)
        return not part or part["subtype"] not in HIDDEN_SUBTYPES
    return [p for p in placements if visible(p)]


def on_plate(project, placements, plate_index):
    plate = next((pl for pl in project["plates"] if pl["index"] == plate_index), None)
    if plate is None:
        raise ValueError(f"no plate {plate_index} in this project")
    wanted = set(plate["objects"])
    return [p for p in placements if (p.item, p.instance) in wanted]


def filament_colour(project, placement):
    """The '#RRGGBB' colour a placed part prints in, or None if unknown.
    A part's extruder of 0 (or none) means it uses its object's extruder."""
    obj = project["objects"].get(placement.item)
    if not obj:
        return None
    part = obj["parts"].get(placement.part) or {}
    extruder = part.get("extruder") or obj["extruder"]
    filaments = project["filaments"]
    if extruder and 1 <= extruder <= len(filaments):
        colour = filaments[extruder - 1]["colour"]
        if isinstance(colour, str) and re.fullmatch(r"#[0-9A-Fa-f]{6}([0-9A-Fa-f]{2})?", colour):
            return colour[:7]
    return None


def measure(meshes, placements):
    """Faces, surface area (mm^2), volume (mm^3) and extents (mm) of everything
    placed, as if concatenated into one mesh — computed per instance to keep
    memory flat."""
    faces, area, volume = 0, 0.0, 0.0
    lo, hi = np.full(3, np.inf), np.full(3, -np.inf)
    for p in placements:
        v, f = meshes[p.key]
        t = p.transform
        w = v @ t[:3, :3].T + t[:3, 3]
        a, b, c = w[f[:, 0]], w[f[:, 1]], w[f[:, 2]]
        cross = np.cross(b - a, c - a)
        faces += len(f)
        area += 0.5 * float(np.linalg.norm(cross, axis=1).sum())
        # A mirroring transform flips winding; count its volume as positive
        # (matching trimesh, which re-winds faces when applying such a transform).
        signed = float(np.einsum("ij,ij->i", a, np.cross(b, c)).sum()) / 6.0
        volume += -signed if np.linalg.det(t[:3, :3]) < 0 else signed
        if len(w):
            lo, hi = np.minimum(lo, w.min(axis=0)), np.maximum(hi, w.max(axis=0))
    extents = (hi - lo) if faces else np.zeros(3)
    return {"faces": faces, "area": area, "volume": abs(volume), "extents": [float(x) for x in extents]}