"""Render a replay stage (worker/usd_scene.py) with Isaac Sim's RTX renderer, one PNG per frame.

Runs inside Isaac Sim's own Python (the container's /isaac-sim/python.sh), not ours:

    /isaac-sim/python.sh worker/isaac_render.py --usd scene.usdc --out frames/ --frames 300 --fps 30

Nothing is simulated here: the stage already holds every body's pose per frame, so each frame only sets the
timeline to that frame's time and renders the camera. Standard library + Isaac Sim + Pillow only.
"""

import argparse
import sys
from pathlib import Path


def main(argv):
    p = argparse.ArgumentParser()
    p.add_argument("--usd", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--frames", type=int, required=True)
    p.add_argument("--fps", type=int, default=30)
    p.add_argument("--width", type=int, default=1280)
    p.add_argument("--height", type=int, default=720)
    p.add_argument("--subframes", type=int, default=4, help="RTX accumulation per frame (quality vs time)")
    args = p.parse_args(argv)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    from isaacsim import SimulationApp

    app = SimulationApp({"headless": True, "width": args.width, "height": args.height})
    try:
        import omni.replicator.core as rep
        import omni.timeline
        import omni.usd
        from PIL import Image

        if not omni.usd.get_context().open_stage(args.usd):
            raise RuntimeError(f"Isaac Sim could not open {args.usd}")
        for _ in range(30):  # let the stage, materials and shaders load
            app.update()
        product = rep.create.render_product("/World/Camera", (args.width, args.height))
        rgb = rep.AnnotatorRegistry.get_annotator("rgb")
        rgb.attach([product])
        timeline = omni.timeline.get_timeline_interface()
        timeline.set_current_time(0.0)
        rep.orchestrator.step(rt_subframes=16, delta_time=0.0, pause_timeline=True)  # warm-up frame
        for i in range(args.frames):
            timeline.set_current_time(i / args.fps)
            rep.orchestrator.step(rt_subframes=args.subframes, delta_time=0.0, pause_timeline=True)
            data = rgb.get_data()
            Image.fromarray(data[:, :, :3]).save(out / f"f{i:05d}.png")
            if i % 30 == 0:
                print(f"[isaac_render] frame {i}/{args.frames}", flush=True)
        rgb.detach()
        product.destroy()
    finally:
        app.close()
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))