"""Bounded synthetic camera/model probe. This is NOT navigation acceptance.

Run with Isaac's python.bat after reserving Worker through the Core job queue.
Outputs stay private in the chosen evidence directory.
"""

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

parser = argparse.ArgumentParser()
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--models", action="store_true")
args = parser.parse_args()
args.output.mkdir(parents=True, exist_ok=False)
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
report = {
    "schema_version": "missioncore.ai-polygon-runtime-probe/v1",
    "started_at": datetime.now(UTC).isoformat(),
    "started_monotonic_ns": time.monotonic_ns(),
    "source": "synthetic-cube-only",
    "navigation_accepted": False,
    "passed": False,
}
app = stack = inference = None
try:
    from isaacsim import SimulationApp

    app = SimulationApp({"headless": True, "multi_gpu": False, "width": 800, "height": 600})
    import numpy as np
    import omni.usd
    from isaacsim.core.experimental.utils import app as app_utils
    from isaacsim.core.simulation_manager import SimulationManager
    from isaacsim.sensors.experimental.rtx import CameraSensor, RtxCamera
    from PIL import Image
    from pxr import Gf, UsdGeom, UsdLux

    stage = omni.usd.get_context().get_stage()
    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
    UsdGeom.SetStageMetersPerUnit(stage, 1.0)
    cube = UsdGeom.Cube.Define(stage, "/World/Cube")
    cube.AddTranslateOp().Set(Gf.Vec3d(0, 3, 0.5))
    cube.CreateDisplayColorAttr([(0.8, 0.15, 0.05)])
    UsdLux.DomeLight.Define(stage, "/World/Light").CreateIntensityAttr(500)
    camera = RtxCamera(
        "/World/Camera",
        tick_rate=10,
        translations=np.array([0.0, 0.0, 0.5]),
        orientations=np.array([1.0, 1.0, 0.0, 0.0]) / np.sqrt(2),
    )
    camera.camera.set_focal_lengths(24.0)
    sensor = CameraSensor(camera, resolution=(600, 800), annotators=["rgb"])
    SimulationManager.setup_simulation(dt=1 / 60, device="cpu")
    app_utils.play()
    app_utils.update_app(steps=12)
    app_utils.pause()
    baseline = SimulationManager.get_num_physics_steps()
    for _ in range(10):
        app.update()
    assert SimulationManager.get_num_physics_steps() == baseline, "Render advanced physics"
    raw, _ = sensor.get_data("rgb")
    assert raw is not None, "RTX camera did not produce a frame"
    rgb = np.ascontiguousarray(raw.numpy()[:, :, :3])
    assert rgb.shape == (600, 800, 3) and rgb.dtype == np.uint8
    assert float(rgb.std()) > 1, "Camera frame is empty/uniform"
    image_path = args.output / "synthetic-camera.png"
    Image.fromarray(rgb).save(image_path)
    report.update(
        frame_sha256=hashlib.sha256(image_path.read_bytes()).hexdigest(),
        frame_shape=list(rgb.shape),
        paused_physics_steps=baseline,
    )
    SimulationManager.step(steps=6)
    assert SimulationManager.get_num_physics_steps() == baseline + 6, "Wrong lockstep increment"
    report["lockstep_physics_steps"] = 6
    if args.models:
        from model_stack import ModelStack

        from k1link.simulation.ai_polygon.inference import ModelInference
        from k1link.simulation.ai_polygon.policy import RoadPolicy

        stack = ModelStack()
        stack.start()
        inference = ModelInference(
            "http://127.0.0.1:18092", Path(stack.profile["labels"]), "http://127.0.0.1:18091"
        )
        inference.ready()
        rows = []
        policy = RoadPolicy(0.3)
        for _ in range(3):
            started = time.monotonic_ns()
            road, boxes = inference.infer(rgb)
            rows.append(
                {
                    "inference_ms": (time.monotonic_ns() - started) / 1e6,
                    "decision": policy.decide(road, boxes).model_dump(),
                }
            )
        report["model_probe"] = rows
    report["passed"] = True
except Exception as exc:
    report["error"] = type(exc).__name__ + ": " + str(exc)
    raise
finally:
    if inference is not None:
        inference.close()
    if stack is not None:
        stack.stop()
    report["finished_at"] = datetime.now(UTC).isoformat()
    (args.output / "report.json").write_text(json.dumps(report, indent=2))
    if app is not None:
        app.close()