"""Rebuild the camera body markers at the new poses, then look through the D435 pair."""
import omni.usd, omni.kit.viewport.utility as vp
import isaacsim.core.experimental.utils.app as app_utils
from pxr import Gf, Usd, UsdGeom, UsdPhysics

stage = omni.usd.get_context().get_stage()
xc = UsdGeom.XformCache()
ROOT = "/World/CameraBodiesSide"
NAMES = ["RealSense_D435_Left", "RealSense_D435_Right",
         "Orbbec_Gemini305_Left", "Orbbec_Gemini305_Right",
         "Orbbec_Gemini345_Left", "Orbbec_Gemini345_Right"]
COLOR = {"RealSense": (0.15, 0.5, 0.95), "Orbbec": (0.95, 0.45, 0.1)}

def cam(name):
    for p in stage.Traverse():
        if p.IsA(UsdGeom.Camera) and p.GetName() == name:
            return p
    raise KeyError(name)

# the old bodies are still at the old poses and would now be both wrong and in the way
old = stage.GetPrimAtPath("/World/CameraBodies")
if old.IsValid():
    UsdGeom.Imageable(old).MakeInvisible()
    for d in Usd.PrimRange(old):
        a = d.GetAttribute("physics:collisionEnabled")
        if a:
            a.Set(False)
    print("old /World/CameraBodies hidden (stale poses, and they occluded the new views)")

UsdGeom.Xform.Define(stage, ROOT)
for n in NAMES:
    c = cam(n)
    M = xc.GetLocalToWorldTransform(c)
    pos = Gf.Vec3d(M.ExtractTranslation())
    fwd = M.TransformDir(Gf.Vec3d(0, 0, -1)); fwd = fwd / (fwd.GetLength() or 1)
    path = f"{ROOT}/{n}"
    if stage.GetPrimAtPath(path).IsValid():
        stage.RemovePrim(path)
    cube = UsdGeom.Cube.Define(stage, path)
    cube.CreateSizeAttr().Set(1.0)
    xf = UsdGeom.Xformable(cube.GetPrim())
    # BEHIND the lens plane: a housing centred on the camera looks into its own inside
    # and the frame comes back black - that failure is documented in this project.
    xf.AddTranslateOp().Set(pos - fwd * 0.045)
    xf.AddScaleOp().Set(Gf.Vec3f(0.05, 0.05, 0.05))
    key = "RealSense" if n.startswith("RealSense") else "Orbbec"
    UsdGeom.Gprim(cube.GetPrim()).CreateDisplayColorAttr().Set([Gf.Vec3f(*COLOR[key])])
print(f"rebuilt {len(NAMES)} body markers under {ROOT}, each offset behind its lens")

w = vp.get_active_viewport()
orig = w.camera_path
shots = []
for n in ("RealSense_D435_Left", "RealSense_D435_Right"):
    w.camera_path = cam(n).GetPath()
    await app_utils.update_app_async(steps=45)
    f = f"/tmp/cam_{n}.png"
    vp.capture_viewport_to_file(w, file_path=f)
    await app_utils.update_app_async(steps=20)
    shots.append(f)
    print("captured", f)
w.camera_path = orig
await app_utils.update_app_async(steps=10)
print("viewport camera restored:", orig)