#!/usr/bin/env python3
"""Build scene/plow_cell.usd from the authored 90_degree.usd.

    python scripts/build_plow_cell.py [path/to/90_degree.usd]

The source is the original authored cell - conveyor art, the Y-split pusher, the plow
(``DiverterEnd``) and its OmniGraph drive script exactly as built. This script does not
regenerate geometry; it only makes the file self-contained inside this repo:

1. **References re-pointed.** The source pulls conveyor art straight off the Omniverse S3
   bucket and the plow meshes from its own folder. Both are re-pointed at the repo's
   ``assets/`` tree so the scene composes offline (``scripts/fetch_assets.py`` fills
   ``assets/conveyors/``).

2. **Baked drive animation stripped.** The pusher's ``PusherSlide`` carries a long
   ``targetPosition.timeSamples`` track. Time samples outrank the attribute default, so
   anything that tries to *control* that drive - the authored script node or Python - is
   overwritten every frame while the timeline runs. The track is dropped; the drive keeps
   its authored gains and limits. (The plow's own ``ArmHinge`` is already clean in
   90_degree.usd; the earlier fixed.usd bakes it too, hence the pattern matches both.)

The authored ``DiverterAnimGraph`` script node is deliberately kept: open the scene, press
Play, and the cell demonstrates itself the way it was built. ``sim/plow_cell.py`` switches
that graph off when you want to drive the plow from code instead.
"""
from __future__ import annotations

import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "scene" / "plow_cell.usd"
DEFAULT_SRC = (Path.home() / "Desktop" / "isaac_sim_project" / "test_isassc" / "90_degree.usd")

S3 = ("https://omniverse-content-production.s3-us-west-2.amazonaws.com/"
      "Assets/Isaac/6.0/Isaac/Props/Conveyors/")

# asset path in the source  ->  path relative to scene/
REFS = {
    f"{S3}ConveyorBelt_A06.usd": "../assets/conveyors/ConveyorBelt_A06.usd",
    f"{S3}ConveyorBelt_A24.usd": "../assets/conveyors/ConveyorBelt_A24.usd",
    "./plow_base.usd": "../assets/plow/plow_base.usd",
    "./plow_arm.usd": "../assets/plow/plow_arm.usd",
}

# Baked drive tracks fight every attempt to control a diverter. In 90_degree.usd only the
# pusher's linear drive carries one (the plow's angular drive is already clean), but the
# earlier fixed.usd bakes the plow too - match both so either source builds the same way.
BAKED = re.compile(r"drive:(linear|angular):physics:targetPosition\.timeSamples")


def usdcat(src: Path, dst: Path) -> None:
    if not shutil.which("usdcat"):
        sys.exit("usdcat not found - it ships with USD / Isaac Sim and is needed to "
                 "convert the binary .usd to text and back")
    subprocess.run(["usdcat", str(src), "-o", str(dst)], check=True)


def strip_baked_track(text: str) -> tuple[str, int]:
    """drop `<drive target>.timeSamples = { ... }` blocks"""
    out, skipping, dropped = [], False, 0
    for line in text.splitlines(keepends=True):
        if not skipping and BAKED.search(line) and line.rstrip().endswith("{"):
            skipping, dropped = True, dropped + 1
            continue
        if skipping:
            if line.strip() == "}":
                skipping = False
            continue
        out.append(line)
    return "".join(out), dropped


def repoint_refs(text: str) -> tuple[str, dict[str, int]]:
    counts = {}
    for old, new in REFS.items():
        n = text.count(f"@{old}@")
        if n:
            text = text.replace(f"@{old}@", f"@{new}@")
        counts[old] = n
    return text, counts


def build(src: Path) -> Path:
    if not src.exists():
        sys.exit(f"source scene not found: {src}")
    OUT.parent.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory() as tmp:
        flat = Path(tmp) / "src.usda"
        usdcat(src, flat)
        text = flat.read_text()

        text, refs = repoint_refs(text)
        text, dropped = strip_baked_track(text)

        missing = [k for k, v in refs.items() if v == 0]
        if missing:
            print("warning: reference not found in source (layout may have changed):")
            for m in missing:
                print(f"  {m}")

        edited = Path(tmp) / "edited.usda"
        edited.write_text(text)
        usdcat(edited, OUT)

    print(f"built {OUT.relative_to(ROOT)} from {src}")
    for old, new in REFS.items():
        print(f"  ref  {refs[old]}x  {Path(old).name:24s} -> {new}")
    print(f"  drop {dropped}x baked drive targetPosition.timeSamples")
    left = re.findall(r"@(https?://[^@]+)@", OUT.read_text(errors="ignore")) if OUT.suffix == ".usda" else []
    if left:
        print(f"  warning: {len(left)} remote reference(s) still present")
    return OUT


if __name__ == "__main__":
    build(Path(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_SRC)