"""Bounded Worker-only drive/braking test, independent of AI and Gaussian mesh."""

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

from isaacsim import SimulationApp

parser = argparse.ArgumentParser()
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()

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 Gf, UsdGeom, UsdPhysics, UsdShade
    from rover_profile import PROFILE, DifferentialDrive, create_rover

    stage = omni.usd.get_context().get_stage()
    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
    UsdGeom.SetStageMetersPerUnit(stage, 1)
    ground = UsdGeom.Cube.Define(stage, "/World/Ground")
    ground.CreateSizeAttr(1)
    ground.AddTranslateOp().Set(Gf.Vec3d(0, 0, -0.1))
    ground.AddScaleOp().Set(Gf.Vec3f(20, 20, 0.2))
    UsdPhysics.CollisionAPI.Apply(ground.GetPrim())
    material = UsdShade.Material.Define(stage, "/World/Material")
    physics = UsdPhysics.MaterialAPI.Apply(material.GetPrim())
    physics.CreateStaticFrictionAttr(0.9)
    physics.CreateDynamicFrictionAttr(0.8)
    UsdShade.MaterialBindingAPI.Apply(ground.GetPrim()).Bind(
        material, UsdShade.Tokens.weakerThanDescendants, "physics"
    )
    create_rover(stage, [0, 0, PROFILE["wheel_radius_m"] + 0.05], 0)
    SimulationManager.setup_simulation(dt=1 / 60, device="cpu")
    omni.timeline.get_timeline_interface().play()
    for _ in range(20):
        app.update()
    robot = Articulation("/World/Rover")
    indices = robot.get_dof_indices(PROFILE["wheel_names"])
    drive = DifferentialDrive(robot)

    def position():
        # Tensor pose is authoritative, independent of USD/Fabric writeback.
        return robot.get_world_poses()[0].numpy().tolist()[0]

    start = position()
    drive.command(0.3)
    SimulationManager.step(steps=300)
    moved = position()
    drive.command(0)
    SimulationManager.step(steps=120)
    stopped = position()
    turns = []

    def heading():
        w, x, y, z = robot.get_world_poses()[1].numpy()[0].tolist()
        return math.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z))

    for direction in (1, -1):
        before, angle = position(), heading()
        drive.command(0, direction * 0.35)
        SimulationManager.step(steps=180)
        turn = math.atan2(math.sin(heading() - angle), math.cos(heading() - angle))
        after = position()
        drive.command(0)
        SimulationManager.step(steps=120)
        turns.append(
            dict(
                command_yaw_rps=direction * 0.35,
                measured_yaw_radians=turn,
                displacement_m=math.dist(before[:2], after[:2]),
                brake_drift_m=math.dist(after, position()),
                passed=direction * turn > 0.5 and math.dist(before[:2], after[:2]) < 0.15,
            )
        )
    report = dict(
        utc=datetime.now(UTC).isoformat(),
        monotonic=time.monotonic(),
        profile_sha256=hashlib.sha256(
            Path(__file__).with_name("rover_profile.py").read_bytes()
        ).hexdigest(),
        profile=PROFILE,
        start=start,
        after_5s=moved,
        after_stop_2s=stopped,
        turns=turns,
        passed=moved[0] - start[0] > 1.0
        and abs(stopped[0] - moved[0]) < 0.15
        and all(turn["passed"] for turn in turns),
        wheel_velocity_rps=robot.get_dof_velocities(dof_indices=indices).numpy().tolist(),
    )
    args.output.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report))
finally:
    app.close()