#!/usr/bin/env python3
"""Flatten the classified objects out of the working scene into assets/items/.

    /home/whatevenif/isaacsim/python.sh scripts/export_item_library.py

The six meshes bundled in `assets/meshes/` are a smoke-test set, two per class. A sorting
run wants the whole catalogue, and the classified objects live inside
`robozon_conveyor_scaled.usd` under `/World/CVObjects`, referencing `.glb` files that are
not in this repo. Each one is therefore *flattened* on export so the result carries its own
geometry and composes anywhere - the recipe in CLAUDE.md, applied in bulk.

Ground truth comes from `categories.json`: `zone` (B/C/D) and `obb_extents_m`, which the
run harness reports predicted dimensions against. Objects larger than the 500 mm incoming
envelope are skipped - they could never reach this conveyor.
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

from pxr import Usd, UsdGeom

ASSETS = Path("/home/dasha/isaac_assets")
SCENE = ASSETS / "robozon_conveyor_scaled.usd"
CATS = ASSETS / "categories.json"
OUT = Path(__file__).resolve().parent.parent / "assets" / "items"
ENVELOPE_MM = 500.0


def main():
    if not SCENE.exists():
        sys.exit(f"{SCENE} not found")
    cats = json.loads(CATS.read_text())
    OUT.mkdir(parents=True, exist_ok=True)

    src = Usd.Stage.Open(str(SCENE))
    root = src.GetPrimAtPath("/World/CVObjects")
    if not root.IsValid():
        sys.exit("/World/CVObjects missing")
    present = {p.GetName() for p in root.GetChildren()}

    manifest, skipped = {}, []
    for name, meta in sorted(cats.items()):
        if name not in present:
            skipped.append((name, "not in scene"))
            continue
        ext_mm = [e * 1000.0 for e in meta["obb_extents_m"]]
        if max(ext_mm) > ENVELOPE_MM:
            skipped.append((name, f"oversize {max(ext_mm):.0f} mm"))
            continue

        ns = Usd.Stage.CreateInMemory()
        item = UsdGeom.Xform.Define(ns, "/Item")
        item.GetPrim().GetReferences().AddReference(str(SCENE), f"/World/CVObjects/{name}")
        ns.SetDefaultPrim(item.GetPrim())
        dst = OUT / f"{name}.usd"
        ns.Flatten().Export(str(dst))

        manifest[name] = dict(zone=meta["zone"],
                              gt_dims_mm=[round(v) for v in sorted(ext_mm, reverse=True)],
                              k_round=round(meta.get("k_round", 0.0), 3),
                              label_ru=meta.get("label_ru", ""))
        print(f"  {name:22s} {meta['zone']}  {manifest[name]['gt_dims_mm']}  "
              f"{dst.stat().st_size / 1024:.0f} KB")

    (OUT / "manifest.json").write_text(json.dumps(manifest, indent=2, ensure_ascii=False))
    from collections import Counter
    print(f"\nexported {len(manifest)} -> {OUT}  "
          f"{dict(Counter(v['zone'] for v in manifest.values()))}")
    if skipped:
        print(f"skipped {len(skipped)}: " + ", ".join(f"{n} ({w})" for n, w in skipped[:8]))


if __name__ == "__main__":
    main()