"""Worker-only contact replay on a retained episode's exact collision asset.

This isolates physical support from perception/planning. It never supplies a
route or collision truth to inference and never writes to the original run.
"""

import argparse
import hashlib
import json
import math
import sys
import time
from datetime import UTC, datetime
from pathlib import Path

from isaacsim import SimulationApp

parser = argparse.ArgumentParser()
parser.add_argument("--run", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--mode", choices=("straight", "replay"), default="straight")
parser.add_argument("--seconds", type=float, default=40)
args = parser.parse_args()
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))

app = SimulationApp({"headless": True, "hide_ui": True})
try:
    import omni.timeline
    import omni.usd
    from isaacsim.core.experimental.prims import Articulation
    from isaacsim.core.simulation_manager import SimulationManager
    from pxr import UsdGeom
    from rover_profile import PROFILE, DifferentialDrive, create_rover
    from terrain import install_terrain

    from k1link.simulation.ai_polygon.mission_policy import inclination

    run = json.loads(args.run.read_text())
    settings = run["world"]["settings"]
    stage = omni.usd.get_context().get_stage()
    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
    UsdGeom.SetStageMetersPerUnit(stage, 1)
    height, terrain = install_terrain(stage, run["terrain_manifest"], run["world"])
    create_rover(
        stage,
        [*settings["spawn_xy"], height],
        settings["heading_degrees"],
        terrain["initial_ground_normal"],
    )
    SimulationManager.setup_simulation(dt=1 / 60, device="cpu")
    omni.timeline.get_timeline_interface().play()
    for _ in range(30):
        app.update()
    robot = Articulation("/World/Rover")
    drive = DifferentialDrive(robot)
    playback = []
    if args.mode == "replay":
        playback = [
            json.loads(x) for x in (args.run.parent / "motion.jsonl").read_text().splitlines()
        ]
    trace, cursor, stopped = [], 0, False
    started = time.monotonic()
    for step in range(round(args.seconds * 60)):
        p, q = robot.get_world_poses()
        p, q = p.numpy()[0].tolist(), q.numpy()[0].tolist()
        pose = p + q[1:] + q[:1]
        tilt = inclination(pose)
        stopped |= tilt >= PROFILE["stop_tilt_degrees"]
        v, w = settings["max_speed_mps"], 0.0
        if playback:
            while (
                cursor + 1 < len(playback)
                and playback[cursor + 1]["simulation_time_ns"] / 1e9 <= step / 60
            ):
                cursor += 1
            v, w = playback[cursor]["applied_speed_mps"], playback[cursor]["applied_yaw_rate_rps"]
        drive.command(0 if stopped else v, 0 if stopped else w)
        if step % 10 == 0:
            trace.append(dict(time_s=step / 60, pose=pose, tilt_degrees=tilt, stopped=stopped))
        SimulationManager.step(steps=1)
    report = dict(
        schema_version="missioncore.scene-contact-qualification/v1",
        utc=datetime.now(UTC).isoformat(),
        monotonic=time.monotonic(),
        elapsed_wall_s=time.monotonic() - started,
        run_sha256=hashlib.sha256(args.run.read_bytes()).hexdigest(),
        mode=args.mode,
        wheel_collision=PROFILE["wheel_collision"],
        seconds=args.seconds,
        profile=PROFILE,
        terrain=terrain,
        trace=trace,
        displacement_m=math.dist(trace[0]["pose"][:2], trace[-1]["pose"][:2]),
        max_tilt_degrees=max(x["tilt_degrees"] for x in trace),
        stopped=stopped,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps({k: v for k, v in report.items() if k not in ("trace", "terrain", "profile")}))
finally:
    app.close()