#!/usr/bin/env python3
"""Give scene/plow_cell.usd an infeed belt, the camera portal and the laser gate.

    python scripts/add_vision_to_plow_cell.py

Purely additive: the conveyors, the Y-split pusher and the plow are left exactly as
authored, including their drives and the DiverterAnimGraph. Nothing existing is edited,
so the cell's kinematics are untouched.

What gets added
  ConveyorTrack_05   a fourth infeed belt, upstream of ConveyorTrack_02. It references the
                     SAME ConveyorBelt_A06.usd with the same scale as the existing tracks,
                     so it is the same conveyor with the same textures, not a lookalike.
                     Goods run -X, so "behind" ConveyorTrack_02 (which spans x -2..0) means
                     x 0..+2.
  CameraMountFrame   the portal, camera bodies and the three rectified stereo pairs, copied
  CameraBodies       verbatim from sorter.usd. They already sit over x=-0.75, which is
  /RigRS             inside ConveyorTrack_02's belt, so no repositioning is needed.
  SortingRig/LaserGate  the through-beam gate at the Y-split pusher (x=-3.74).

Re-running is safe: existing prims are replaced rather than duplicated.
"""
from __future__ import annotations

import shutil
import sys
from pathlib import Path

from pxr import Gf, Sdf, Usd, UsdGeom

ROOT = Path(__file__).resolve().parent.parent
CELL = ROOT / "scene" / "plow_cell.usd"
SORTER = ROOT / "scene" / "sorter.usd"

# copied from sorter.usd unchanged - they are already aligned with this belt
FROM_SORTER = [
    "/World/CameraMountFrame",
    "/World/CameraBodies",
    "/RigRS",
    "/World/SortingRig",          # carries LaserGate (and its materials)
]

# sorter.usd's SortingRig carries a plain box SpawnBelt at x 0..2.6. The real conveyor
# added below occupies the same lane, so the box and its rails/legs are dropped after the
# copy - otherwise two belts sit inside each other. LaserGate, the bin and the materials stay.
DROP_AFTER_COPY = [
    "/World/SortingRig/SpawnBelt",
    "/World/SortingRig/SpawnRail_p", "/World/SortingRig/SpawnRail_n",
    "/World/SortingRig/Spawn_Leg0", "/World/SortingRig/Spawn_Leg1",
    "/World/SortingRig/Spawn_Leg2", "/World/SortingRig/Spawn_Leg3",
]

INFEED_PRIM = "/World/ConveyorTrack_05"
INFEED_ASSET = "../assets/conveyors/ConveyorBelt_A06.usd"
# ConveyorTrack_02 sits at translate x=-2 and spans x -2..0, so the asset occupies
# [tx, tx+2]. Upstream of it is therefore tx=0 -> x 0..+2.
INFEED_TRANSLATE = Gf.Vec3d(0.0, 0.0, 0.0)
INFEED_SCALE = Gf.Vec3d(1.0, 0.5, 1.0)          # identical to the other tracks


def _drop(layer: Sdf.Layer, path: str):
    """remove a prim spec if present, so the script is idempotent"""
    spec = layer.GetPrimAtPath(path)
    if not spec:
        return False
    parent = layer.GetPrimAtPath(str(Sdf.Path(path).GetParentPath())) or layer.pseudoRoot
    name = Sdf.Path(path).name
    if name in parent.nameChildren:
        del parent.nameChildren[name]
        return True
    return False


def add_infeed(layer: Sdf.Layer):
    """a fourth conveyor upstream of ConveyorTrack_02, same asset and scale"""
    _drop(layer, INFEED_PRIM)
    spec = Sdf.CreatePrimInLayer(layer, INFEED_PRIM)
    spec.specifier = Sdf.SpecifierDef
    spec.typeName = "Xform"
    spec.referenceList.prependedItems.append(Sdf.Reference(INFEED_ASSET))
    for name, value, vtype in (
            ("xformOp:translate", INFEED_TRANSLATE, Sdf.ValueTypeNames.Double3),
            ("xformOp:scale", INFEED_SCALE, Sdf.ValueTypeNames.Double3)):
        attr = Sdf.AttributeSpec(spec, name, vtype)
        attr.default = value
    order = Sdf.AttributeSpec(spec, "xformOpOrder", Sdf.ValueTypeNames.TokenArray)
    order.default = ["xformOp:translate", "xformOp:scale"]
    return INFEED_PRIM


def copy_from_sorter(layer: Sdf.Layer, src: Sdf.Layer):
    copied = []
    for path in FROM_SORTER:
        if not src.GetPrimAtPath(path):
            print(f"  skip {path} - not in sorter.usd")
            continue
        _drop(layer, path)
        if Sdf.CopySpec(src, Sdf.Path(path), layer, Sdf.Path(path)):
            copied.append(path)
    return copied


def main():
    if not CELL.exists():
        print(f"{CELL} not found - build it with scripts/build_plow_cell.py")
        return 1
    if not SORTER.exists():
        print(f"{SORTER} not found - the camera stand is copied from it")
        return 1

    backup = CELL.with_suffix(".usd.bak")
    if backup.exists():
        print(f"backup {backup.name} already exists - keeping the pre-vision copy")
    else:
        shutil.copy(CELL, backup)
        print(f"backup -> {backup.name}")

    layer = Sdf.Layer.FindOrOpen(str(CELL))
    src = Sdf.Layer.FindOrOpen(str(SORTER))

    infeed = add_infeed(layer)
    print(f"added {infeed}  (references {INFEED_ASSET}, translate {tuple(INFEED_TRANSLATE)})")

    for path in copy_from_sorter(layer, src):
        print(f"copied {path}")

    for path in DROP_AFTER_COPY:
        if _drop(layer, path):
            print(f"dropped {path} (superseded by the real conveyor)")

    layer.Save()
    print(f"saved {CELL}")

    # verify by composing the result
    stage = Usd.Stage.Open(str(CELL))
    cache = UsdGeom.BBoxCache(0, ["default", "render"], useExtentsHint=True)
    print("\nverification:")
    ok = True
    for path in [f"{INFEED_PRIM}/Belt", "/World/ConveyorTrack_02/Belt",
                 "/World/CameraMountFrame", "/World/SortingRig/LaserGate"]:
        prim = stage.GetPrimAtPath(path)
        if not prim.IsValid():
            print(f"  MISSING {path}")
            ok = False
            continue
        r = cache.ComputeWorldBound(prim).ComputeAlignedRange()
        if r.IsEmpty():
            print(f"  EMPTY   {path}")
            ok = False
            continue
        mn, mx = r.GetMin(), r.GetMax()
        print(f"  ok      {path}: x[{mn[0]:7.3f}..{mx[0]:7.3f}] y[{mn[1]:6.3f}..{mx[1]:6.3f}] "
              f"top_z={mx[2]:.3f}")
    cams = stage.GetPrimAtPath("/RigRS")
    n = len(cams.GetChildren()) if cams.IsValid() else 0
    print(f"  {'ok     ' if n == 6 else 'PROBLEM'} /RigRS: {n} cameras")
    return 0 if ok and n == 6 else 1


if __name__ == "__main__":
    sys.exit(main())