"""Native Worker-local physics/render/camera loop, independent of Core and AI latency."""

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

from local_state import StateChannel, read_json, write_json
from motion_control import CONTROL_PROFILE, DriveEnvelope
from realtime_ai import LatestInference
from rover_profile import PROFILE, DifferentialDrive, create_rover

parser = argparse.ArgumentParser()
parser.add_argument("--run", type=Path, required=True)
parser.add_argument("--source", type=Path, required=True)
parser.add_argument("--stream-address", required=True)
args, _ = parser.parse_known_args()
root = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(root / "src"))
run = read_json(args.run)
settings = run["world"]["settings"]
directory = args.run.parent
channel = StateChannel(directory)
result = {"outcome": "failed", "message": "Симуляция не завершена."}
app = ai = None
physics_callbacks = []


def startup_phase(phase):
    record = {"phase": phase, "utc": datetime.now(UTC).isoformat(), "monotonic": time.monotonic()}
    write_json(directory / "startup.json", record)
    print(json.dumps({"startup": record}), flush=True)


try:
    startup_phase("native-runtime")
    from isaacsim import SimulationApp

    app = SimulationApp(
        {
            "headless": True,
            "hide_ui": True,
            "multi_gpu": False,
            "width": 1280,
            "height": 720,
            "window_width": 1280,
            "window_height": 720,
            "renderer": "RaytracedLighting",
            "display_options": 0,
            "extra_args": [
                "--/app/player/useFixedTimeStepping=false",
                "--/app/runLoops/main/manualModeEnabled=false",
                "--/exts/isaacsim.core.throttling/enable_manualmode=false",
            ],
        }
    )
    import carb.settings
    import numpy as np
    import omni.kit.app
    import omni.replicator.core as rep
    import omni.timeline
    import omni.usd
    from isaacsim.core.experimental.prims import Articulation
    from isaacsim.core.experimental.utils.app import enable_extension
    from isaacsim.core.simulation_manager import SimulationEvent, SimulationManager
    from navigation.footprint import FRAME_DEADLINE_SECONDS
    from navigation_client import ComposedInference
    from omni.kit.loop import _loop as omni_loop
    from omni.kit.viewport.utility import get_active_viewport
    from pxr import Gf, UsdGeom, UsdLux
    from terrain import RangeSensor, install_terrain

    from k1link.simulation.ai_polygon.mission_policy import inclination

    startup_phase("stream-and-gaussian-runtime")
    config = carb.settings.get_settings()
    prefix = "/exts/omni.kit.livestream.app/primaryStream/"
    for key, value in {
        "publicIp": args.stream_address,
        "signalPort": 49100,
        "streamPort": 47998,
        "targetFps": 30,
        "enableEventTracing": False,
    }.items():
        config.set(prefix + key, value)
    enable_extension("omni.kit.livestream.app")
    enable_extension("omni.kit.converter.gsplat")
    from usd_convert_gsplat import read_ply, write_gaussian_splat_usd

    converted = args.source.with_suffix(".usd")
    if not converted.exists():
        temporary = converted.with_suffix(".part.usd")
        write_gaussian_splat_usd(
            read_ply(str(args.source)),
            str(temporary),
            source_file=str(args.source),
            prim_name="Gaussians",
            up_axis="Z",
        )
        temporary.replace(converted)
    stage = omni.usd.get_context().get_stage()
    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
    UsdGeom.SetStageMetersPerUnit(stage, 1)
    visual = UsdGeom.Xform.Define(stage, "/World/Scan")
    visual.AddRotateXYZOp().Set(Gf.Vec3f(*settings["rotation_degrees"]))
    visual.AddScaleOp().Set(Gf.Vec3f(settings["meters_per_unit"]))
    stage.DefinePrim("/World/Scan/Gaussians").GetReferences().AddReference(str(converted))
    startup_phase("terrain")
    wheel_center_z, terrain_record = install_terrain(stage, run["terrain_manifest"], run["world"])
    write_json(directory / "terrain.json", terrain_record)
    UsdLux.DomeLight.Define(stage, "/World/Light").CreateIntensityAttr(500)
    create_rover(
        stage,
        [*settings["spawn_xy"], wheel_center_z],
        settings["heading_degrees"],
        terrain_record["initial_ground_normal"],
    )
    write_json(directory / "robot-profile.json", PROFILE)
    write_json(directory / "motion-control.json", CONTROL_PROFILE)
    startup_phase("cameras")

    def camera(path, aspect, focal_length=24):
        value = UsdGeom.Camera.Define(stage, path)
        value.CreateFocalLengthAttr(focal_length)
        value.CreateHorizontalApertureAttr(36)
        value.CreateVerticalApertureAttr(36 / aspect)
        value.CreateClippingRangeAttr(Gf.Vec2f(0.05, 1000))
        return value, value.AddTranslateOp(), value.AddOrientOp()

    sensor, sensor_pos, sensor_rot = camera(
        "/World/Sensor", 4 / 3, PROFILE["camera_focal_length_mm"]
    )
    focal_pixels = (
        800 * PROFILE["camera_focal_length_mm"] / PROFILE["camera_horizontal_aperture_mm"]
    )
    observer, observer_pos, observer_rot = camera("/World/Observer", 16 / 9)
    product = rep.create.render_product(str(sensor.GetPath()), (800, 600))
    annotator = rep.AnnotatorRegistry.get_annotator("rgb")
    annotator.attach(product)
    config.set("/omni/replicator/captureOnPlay", True)
    viewport = get_active_viewport()
    if viewport is None:
        raise RuntimeError("Streaming viewport is unavailable")
    viewport.camera_path = str(observer.GetPath())
    viewport.set_texture_resolution((1280, 720))
    # Initialise the final camera before the first rendered frames. Starting
    # WebRTC with an uninitialised observer and an already paused timeline can
    # leave the native stream without its first image/offer.
    initial = Gf.Vec3d(*settings["spawn_xy"], wheel_center_z + 0.12)
    heading = math.radians(settings["heading_degrees"])
    initial_eye = initial + Gf.Vec3d(-math.cos(heading) * 2.5, -math.sin(heading) * 2.5, 1.5)
    observer_pos.Set(initial_eye)
    observer_rot.Set(
        Gf.Quatf(
            Gf.Matrix4d()
            .SetLookAt(initial_eye, initial + Gf.Vec3d(0.1, 0, 0.2), Gf.Vec3d(0, 0, 1))
            .GetInverse()
            .ExtractRotationQuat()
        )
    )
    startup_phase("physics-setup")
    SimulationManager.setup_simulation(dt=1 / 60, device="cpu")
    timeline = omni.timeline.get_timeline_interface()
    # Render rate and physics rate are separate. Physics steps follow elapsed time.
    config.set("/app/player/useFixedTimeStepping", False)
    config.set("/exts/isaacsim.core.throttling/enable_manualmode", False)
    omni_loop.acquire_loop_interface().set_manual_mode(False)
    timeline.set_play_every_frame(False)
    config.set("/persistent/simulation/minFrameRate", 1)
    config.set("/app/runLoops/main/rateLimitEnabled", True)
    config.set("/app/runLoops/main/rateLimitFrequency", 30)
    startup_phase("initial-render-and-contact")
    timeline.play()
    for _ in range(30):
        app.update()
    articulation = Articulation("/World/Rover")
    drive = DifferentialDrive(articulation)
    wheel_indices = articulation.get_dof_indices(PROFILE["wheel_names"])

    def body_transform():
        positions, orientations = articulation.get_world_poses()
        position, orientation = positions.numpy()[0], orientations.numpy()[0]
        transform = Gf.Matrix4d(1)
        transform.SetRotate(Gf.Quatd(float(orientation[0]), Gf.Vec3d(*map(float, orientation[1:]))))
        transform.SetTranslateOnly(Gf.Vec3d(*map(float, position)))
        return transform

    settled = body_transform().ExtractTranslation()
    write_json(
        directory / "initial-contact.json",
        {
            "pose": list(settled),
            "expected_xy": settings["spawn_xy"],
            "spawn_contact_z": terrain_record["initial_contact_z"],
        },
    )
    if math.hypot(settled[0] - settings["spawn_xy"][0], settled[1] - settings["spawn_xy"][1]) > 0.2:
        raise RuntimeError("Rover start is unstable on this reconstructed surface")
    timeline.pause()
    app.update()
    startup_phase("ready")
    baseline = SimulationManager.get_num_physics_steps()
    range_sensor = RangeSensor()
    from k1link.simulation.ai_polygon.mission_policy import WaypointMission

    # The episode owns the mission. Reconnecting a failed model client must not
    # rewind its route cursor, recovery budget or latched terminal condition.
    mission = WaypointMission(settings.get("route_xy", []))
    ai = LatestInference(
        lambda: ComposedInference(
            root / "simulation/ai-polygon", run, directory / "camera", mission=mission
        ),
        None,
        directory / "camera",
    )
    ai.start()
    control = {"control": "pause", "control_sequence": 0, "camera": "follow"}
    start = last_report = last_sensor = last_control = time.monotonic()
    render_frames = sensor_frames = sequence = 0
    prior_metrics = (start, 0, 0, 0, baseline)
    speed = 0.0
    had_ai = False
    unstable = False
    terminal_at = None
    envelope = DriveEnvelope(track_width=PROFILE["track_width_m"])
    state = {"velocity": 0.0, "yaw_rate": 0.0, "stop_reason": "paused", "decision": None}
    physics_rows = []

    def control_step(dt, _context):
        velocity, yaw_rate, reason, decision = ai.command(
            time.monotonic(), frame_deadline=FRAME_DEADLINE_SECONDS
        )
        terminal = unstable or (
            decision and decision["decision"]["reason"] in {"unstable", "stuck", "goal-reached"}
        )
        velocity, yaw_rate = envelope.step(
            velocity, yaw_rate, dt, stop=bool(terminal) or reason != "none"
        )
        drive.command(velocity, yaw_rate)
        state.update(
            velocity=velocity,
            yaw_rate=yaw_rate,
            stop_reason="unstable" if unstable else reason,
            decision=decision,
        )

    def update_cameras():
        transform = body_transform()
        position = transform.ExtractTranslation()
        forward = transform.TransformDir(Gf.Vec3d(1, 0, 0)).GetNormalized()
        yaw = math.atan2(forward[1], forward[0])
        # Both sensor position and attitude follow the articulated body.
        # Mounting height is measured from nominal wheel contact.
        eye = transform.Transform(
            Gf.Vec3d(
                PROFILE["camera_forward_m"],
                0,
                settings["camera_height_m"] - PROFILE["body_contact_height_m"],
            )
        )
        pitch = math.radians(PROFILE["camera_pitch_degrees"])
        camera_axes = [
            transform.TransformDir(Gf.Vec3d(*axis)).GetNormalized()
            for axis in (
                (math.cos(pitch), 0, math.sin(pitch)),
                (0, 1, 0),
                (-math.sin(pitch), 0, math.cos(pitch)),
            )
        ]

        def look_at(pos, rotation, origin, target, up=None):
            if up is None:
                up = Gf.Vec3d(0, 0, 1)
            pos.Set(origin)
            rotation.Set(
                Gf.Quatf(
                    Gf.Matrix4d().SetLookAt(origin, target, up).GetInverse().ExtractRotationQuat()
                )
            )

        look_at(
            sensor_pos,
            sensor_rot,
            eye,
            eye + camera_axes[0],
            camera_axes[2],
        )
        mode = control.get("camera", "follow")
        if mode == "camera":
            viewport.camera_path = str(sensor.GetPath())
        else:
            viewport.camera_path = str(observer.GetPath())
            offset = (
                Gf.Vec3d(-math.cos(yaw) * 2.5, -math.sin(yaw) * 2.5, 1.5)
                if mode == "follow"
                else Gf.Vec3d(-3, -3, 5)
            )
            look_at(observer_pos, observer_rot, position + offset, position + Gf.Vec3d(0.1, 0, 0.2))
        linear, angular = articulation.get_velocities()
        q = transform.ExtractRotationQuat()
        state.update(
            transform=transform,
            position=position,
            yaw=yaw,
            eye=eye,
            camera_axes=camera_axes,
            mode=mode,
            pose=list(position) + list(q.GetImaginary()) + [q.GetReal()],
            speed=float(np.linalg.norm(linear.numpy()[0, :2])),
            body_velocity=linear.numpy()[0].tolist(),
            body_angular_velocity=angular.numpy()[0].tolist(),
        )

    def physics_step(dt, _context):
        global unstable
        update_cameras()
        unstable |= inclination(state["pose"]) >= PROFILE["stop_tilt_degrees"]
        physics_rows.append(
            {
                "monotonic": time.monotonic(),
                "dt": dt,
                "pose": state["pose"],
                "velocity": state["body_velocity"],
                "angular_velocity": state["body_angular_velocity"],
                "command": [state["velocity"], state["yaw_rate"]],
                "stop_reason": state["stop_reason"],
            }
        )

    callback_errors = []

    def guarded(callback):
        def invoke(dt, context):
            try:
                callback(dt, context)
            except Exception:
                # Native event dispatch logs Python exceptions and continues.
                # A failed control/camera callback must instead end this run.
                callback_errors.append(traceback.format_exc())
                ai.enable(False)
                drive.command(0.0, 0.0)

        return invoke

    update_cameras()
    physics_callbacks.extend(
        [
            SimulationManager.register_callback(
                guarded(control_step), SimulationEvent.PHYSICS_PRE_STEP
            ),
            SimulationManager.register_callback(
                guarded(physics_step), SimulationEvent.PHYSICS_POST_STEP
            ),
        ]
    )
    with (
        (directory / "motion.jsonl").open("a", encoding="utf-8") as motion,
        (directory / "physics-motion.jsonl").open("a", encoding="utf-8") as physics_motion,
    ):
        while app.is_running():
            now = time.monotonic()
            if now - last_control >= 0.1:
                control = channel.read("control", control)
                last_control = now
            if control["control"] == "stop":
                result = {"outcome": "stopped", "message": "Симуляция и inference остановлены."}
                break
            if now - start >= run["request"]["duration_seconds"]:
                result = {"outcome": "completed", "message": "Время прогона завершено."}
                break
            playing = control["control"] == "play"
            ai.enable(playing)
            had_ai |= playing
            if playing and not timeline.is_playing():
                timeline.play()
            elif not playing and timeline.is_playing():
                timeline.pause()
            velocity, yaw_rate, stop_reason, decision = ai.command(
                now, frame_deadline=FRAME_DEADLINE_SECONDS
            )
            transform = body_transform()
            q = transform.ExtractRotationQuat()
            pose = list(transform.ExtractTranslation()) + list(q.GetImaginary()) + [q.GetReal()]
            unstable |= inclination(pose) >= PROFILE["stop_tilt_degrees"]
            if unstable:
                velocity, yaw_rate, stop_reason = 0.0, 0.0, "unstable"
            terminal_reason = (
                "unstable" if unstable else (decision["decision"]["reason"] if decision else None)
            )
            if terminal_reason in {"unstable", "stuck", "goal-reached"}:
                velocity, yaw_rate = 0.0, 0.0
                terminal_at = now if terminal_at is None else terminal_at
                if now - terminal_at >= 2:
                    result = {
                        "outcome": "completed" if terminal_reason == "goal-reached" else "failed",
                        "message": {
                            "goal-reached": "Ровер достиг цели маршрута.",
                            "stuck": "Ровер не нашёл проезд после трёх попыток. Прогон завершён.",
                            "unstable": "Прогон завершён из-за опасного наклона ровера.",
                        }[terminal_reason],
                    }
                    break
            # The physics tensor API applies live drive targets in radians/s;
            # USD authoring attributes are only the initial scene configuration.
            if not playing:
                drive.command(*envelope.step(0.0, 0.0, 0.0, stop=True))
                update_cameras()
                state.update(
                    velocity=0.0, yaw_rate=0.0, stop_reason="paused", speed=0.0, decision=decision
                )
            app.update()  # Never waits for inference, a network request, or an operator ACK.
            if callback_errors:
                raise RuntimeError(callback_errors[0])
            render_frames += 1
            captured_at = time.monotonic()
            transform, position = state["transform"], state["position"]
            eye, camera_axes = state["eye"], state["camera_axes"]
            pose, yaw, speed, mode = state["pose"], state["yaw"], state["speed"], state["mode"]
            velocity, yaw_rate = state["velocity"], state["yaw_rate"]
            stop_reason, decision = state["stop_reason"], state["decision"]
            if physics_rows:
                physics_motion.writelines(
                    json.dumps(row, allow_nan=False) + "\n" for row in physics_rows
                )
                physics_rows.clear()
                physics_motion.flush()
            physics = SimulationManager.get_num_physics_steps()
            sim_ns = round((physics - baseline) * 1e9 / 60)
            if playing and captured_at - last_sensor >= 1 / 5:
                rgb = annotator.get_data()
                if isinstance(rgb, np.ndarray) and rgb.shape[:2] == (600, 800):
                    sensor_frames += 1
                    rear_eye = transform.Transform(
                        Gf.Vec3d(
                            PROFILE["rear_range_forward_m"],
                            0,
                            settings["camera_height_m"] - PROFILE["body_contact_height_m"],
                        )
                    )
                    points = range_sensor.capture(eye, transform, rear_eye)
                    quaternion = transform.ExtractRotationQuat()
                    imaginary = quaternion.GetImaginary()
                    # Columns are the camera's forward/left/up axes in world.
                    observation = {
                        "simulation_time_ns": sim_ns,
                        "points": points,
                        "pose": [float(v) for v in position]
                        + [float(v) for v in imaginary]
                        + [float(quaternion.GetReal())],
                        "calibration": {
                            "origin": list(eye),
                            "rotation": np.array(camera_axes).T.reshape(-1).tolist(),
                            "intrinsics": [focal_pixels, focal_pixels, 400.0, 300.0],
                            "body_contact_height_m": PROFILE["body_contact_height_m"],
                            "range_origins": [list(eye), list(rear_eye)],
                        },
                    }
                    ai.submit(
                        np.array(rgb[:, :, :3], copy=True, order="C"),
                        sensor_frames,
                        captured_at,
                        sim_ns,
                        observation,
                    )
                last_sensor = captured_at
            if captured_at - last_report >= 0.25:
                prior_at, prior_render, prior_sensor, prior_ai, prior_physics = prior_metrics
                elapsed = captured_at - prior_at
                report = dict(
                    sequence=sequence,
                    control_sequence=control["control_sequence"],
                    state="running" if playing else "paused" if had_ai else "ready",
                    phase="models" if playing and not ai.ever_ready else "running",
                    simulation_time_ns=sim_ns,
                    wall_elapsed_seconds=captured_at - start,
                    physics_steps=physics - baseline,
                    render_frames=render_frames,
                    sensor_frames=sensor_frames,
                    inference_count=ai.count,
                    dropped_frames=ai.dropped,
                    rtf=(physics - prior_physics) / 60 / elapsed,
                    render_fps=(render_frames - prior_render) / elapsed,
                    sensor_fps=(sensor_frames - prior_sensor) / elapsed,
                    ai_hz=(ai.count - prior_ai) / elapsed,
                    inference_ms=decision["inference_ms"] if decision else None,
                    frame_age_ms=(captured_at - decision["captured_at"]) * 1000
                    if decision
                    else None,
                    command_age_ms=(captured_at - decision["completed_at"]) * 1000
                    if decision
                    else None,
                    pose_xy=[float(position[0]), float(position[1])],
                    pose_yaw=yaw,
                    speed_mps=speed,
                    applied_speed_mps=velocity,
                    applied_yaw_rate_rps=yaw_rate,
                    decision=decision["decision"] if decision else None,
                    stop_reason=stop_reason,
                    ai_ready=ai.ready,
                    stream_ready=True,
                    camera=mode,
                )
                channel.write("snapshot", report)
                motion.write(
                    json.dumps(
                        {
                            **report,
                            "monotonic": captured_at,
                            "pose_xyz_xyzw": pose,
                            "tilt_degrees": inclination(pose),
                            "wheel_velocity_rps": articulation.get_dof_velocities(
                                dof_indices=wheel_indices
                            )
                            .numpy()
                            .tolist()[0],
                            "usd_position_xyz": list(
                                UsdGeom.XformCache()
                                .GetLocalToWorldTransform(
                                    stage.GetPrimAtPath("/World/Rover/chassis")
                                )
                                .ExtractTranslation()
                            ),
                        },
                        allow_nan=False,
                    )
                    + "\n"
                )
                motion.flush()
                last_report = captured_at
                sequence += 1
                if elapsed >= 1:
                    prior_metrics = (captured_at, render_frames, sensor_frames, ai.count, physics)
                if shutil.disk_usage(directory).free < 1024**3:
                    raise RuntimeError("Worker evidence disk is full")
except Exception as exc:
    message = str(exc)
    placement_failure = any(
        text in message
        for text in (
            "Configured spawn",
            "Rover start",
            "reconstructed ground",
            "reconstructed support",
        )
    )
    result = {
        "outcome": "failed",
        "message": "Не удалось устойчиво разместить ровер на грунте. Измените положение старта."
        if placement_failure
        else "Симуляция прервана из-за ошибки на Worker. Журнал сохранён.",
    }
    traceback.print_exc()
finally:
    for callback_id in physics_callbacks:
        SimulationManager.deregister_callback(callback_id)
    channel.close()
    if ai is not None:
        ai.close()
    write_json(directory / "result.json", result)
    if app is not None:
        app.close()