"""Streaming 3MF reader shared by convert_3mf.py and fingerprint.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.
"""
import zipfile
from array import array

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


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 _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: [(key, 4x4 transform)] — one per placed instance, mapping 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):
            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((key, transform))
            for cpath, coid, m in comps:
                walk(cpath.lstrip("/") if cpath else name, coid, transform @ m, depth + 1)

        for oid, m in items:
            walk(root, oid, scale @ m, 0)
    return meshes, placements


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 key, t in placements:
        v, f = meshes[key]
        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]}