#!/usr/bin/env python3
"""Replace one object's mesh inside a .3mf file with an STL, writing a new .3mf.

Preserves all other objects, build items, metadata, and auxiliary assets
(thumbnails, etc.) from the source 3mf. Only the mesh data of the target
object id is substituted.
"""
from __future__ import annotations

import argparse
import os
import re
import shutil
import struct
import sys
import tempfile
import zipfile
from pathlib import Path


def read_stl(path: Path) -> tuple[list[tuple[float, float, float]], list[tuple[int, int, int]]]:
    data = path.read_bytes()
    is_ascii = data[:6].lstrip().startswith(b"solid") and b"facet" in data[:4096]

    pts: dict[tuple[float, float, float], int] = {}
    vlist: list[tuple[float, float, float]] = []
    tris: list[tuple[int, int, int]] = []

    def intern(v: tuple[float, float, float]) -> int:
        key = (round(v[0], 6), round(v[1], 6), round(v[2], 6))
        if key not in pts:
            pts[key] = len(vlist)
            vlist.append(key)
        return pts[key]

    if is_ascii:
        raw = re.findall(r"vertex\s+(\S+)\s+(\S+)\s+(\S+)", data.decode("utf-8", errors="ignore"))
        for i in range(0, len(raw), 3):
            face = tuple(
                intern((float(raw[i + k][0]), float(raw[i + k][1]), float(raw[i + k][2])))
                for k in range(3)
            )
            tris.append(face)
    else:
        n = struct.unpack_from("<I", data, 80)[0]
        offset = 84
        for _ in range(n):
            tri = struct.unpack_from("<12fH", data, offset)
            offset += 50
            face = tuple(
                intern((tri[3 + 3 * k], tri[4 + 3 * k], tri[5 + 3 * k]))
                for k in range(3)
            )
            tris.append(face)

    return vlist, tris


def mesh_xml(
    verts: list[tuple[float, float, float]],
    tris: list[tuple[int, int, int]],
    indent: str = "   ",
) -> str:
    lines = [f"{indent}<mesh>", f"{indent} <vertices>"]
    lines.extend(
        f'{indent}  <vertex x="{x}" y="{y}" z="{z}" />' for x, y, z in verts
    )
    lines.append(f"{indent} </vertices>")
    lines.append(f"{indent} <triangles>")
    lines.extend(
        f'{indent}  <triangle v1="{a}" v2="{b}" v3="{c}" />' for a, b, c in tris
    )
    lines.append(f"{indent} </triangles>")
    lines.append(f"{indent}</mesh>")
    return "\n".join(lines)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--source", type=Path, required=True, help="original .3mf file")
    ap.add_argument("--object-id", required=True, help="object id whose mesh is replaced")
    ap.add_argument("--replacement", type=Path, required=True, help=".stl with the new mesh")
    ap.add_argument("--output", "-o", type=Path, required=True, help="output .3mf path")
    ap.add_argument(
        "--model-path",
        default="3D/3dmodel.model",
        help="path of the model XML inside the 3mf (default: %(default)s)",
    )
    args = ap.parse_args()

    verts, tris = read_stl(args.replacement)
    new_mesh = mesh_xml(verts, tris)

    with tempfile.TemporaryDirectory() as tmp:
        staged = Path(tmp) / "staged"
        with zipfile.ZipFile(args.source) as z:
            z.extractall(staged)

        model_file = staged / args.model_path
        xml = model_file.read_text()
        pattern = re.compile(
            rf'(<object id="{re.escape(args.object_id)}"[^>]*>)(.*?)(</object>)',
            re.DOTALL,
        )
        if not pattern.search(xml):
            print(f"object id={args.object_id!r} not found", file=sys.stderr)
            return 1
        new_xml = pattern.sub(
            lambda m: m.group(1) + "\n" + new_mesh + "\n  " + m.group(3), xml
        )
        model_file.write_text(new_xml)

        args.output.parent.mkdir(parents=True, exist_ok=True)
        if args.output.exists():
            args.output.unlink()
        with zipfile.ZipFile(args.output, "w", zipfile.ZIP_DEFLATED) as z:
            ct = staged / "[Content_Types].xml"
            if ct.exists():
                z.write(ct, "[Content_Types].xml")
            for root, _dirs, files in os.walk(staged):
                for fn in files:
                    p = Path(root) / fn
                    rel = p.relative_to(staged).as_posix()
                    if rel == "[Content_Types].xml":
                        continue
                    z.write(p, rel)

    print(f"wrote {args.output} ({len(verts)} verts, {len(tris)} tris in object {args.object_id})")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())