#!/usr/bin/env python3
"""Offline H2 artifact/mapping/policy/CLoSD preflight; never opens robot topics."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys
import time

import numpy as np
import torch


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from runtime.configuration import load_settings  # noqa: E402
from runtime.embodiments import load_embodiment  # noqa: E402
from runtime.h2_contract import (  # noqa: E402
    EXPECTED_HW_TO_POLICY,
    EXPECTED_POLICY_TO_HW,
    HARDWARE_JOINT_NAMES,
    H2_JOINT_MAP_REVISION,
    H2Contract,
    H2_DOF,
)
from runtime.h2_shm_bridge import BRIDGE_SIZE, SHM_VERSION  # noqa: E402
from runtime.observation import build_observation, expected_obs_dim  # noqa: E402
from runtime.policy import OnnxPolicy  # noqa: E402
from runtime.provenance import verify_artifacts  # noqa: E402
from runtime.deployment_versions import DEPLOYMENTS  # noqa: E402
from runtime.real_reference import RealReferenceState  # noqa: E402
from runtime.reference import ReferenceGenerator, load_frozen_pose  # noqa: E402


def run(config: Path, *, policy_only: bool, deployment_version: str | None = None) -> dict:
    settings = load_settings(config, deployment_version=deployment_version)
    verified = verify_artifacts()
    spec = load_embodiment("h2")
    policy = OnnxPolicy(str(settings.policy_path), device=settings.policy_device)
    contract = H2Contract.from_policy(policy, spec.mjcf_path())
    frozen = load_frozen_pose(spec, list(contract.policy_names))
    ready_hw = contract.ready_hw(frozen.joint_pos)

    probe_hw = np.arange(H2_DOF, dtype=np.float64)
    roundtrip = contract.hw_from_policy(contract.policy_from_hw(probe_hw))
    if not np.array_equal(probe_hw, roundtrip):
        raise RuntimeError("H2 mapping round-trip failed")
    if tuple(contract.hw_to_policy.tolist()) != EXPECTED_HW_TO_POLICY:
        raise RuntimeError("H2 hardware-to-policy mapping tripwire failed")
    if tuple(contract.policy_to_hw.tolist()) != EXPECTED_POLICY_TO_HW:
        raise RuntimeError("H2 policy-to-hardware mapping tripwire failed")
    if settings.hardware_joint_map_revision != H2_JOINT_MAP_REVISION:
        raise RuntimeError(
            "configured H2 joint-map revision does not match the audited executable: "
            f"{settings.hardware_joint_map_revision!r} != {H2_JOINT_MAP_REVISION!r}"
        )
    critical_joint_slots = {
        4: ("left_ankle_roll_joint", 14),
        5: ("left_ankle_pitch_joint", 19),
        10: ("right_ankle_roll_joint", 15),
        11: ("right_ankle_pitch_joint", 20),
        12: ("waist_yaw_joint", 2),
        13: ("waist_roll_joint", 5),
        14: ("waist_pitch_joint", 8),
        19: ("left_wrist_roll_joint", 25),
        20: ("left_wrist_pitch_joint", 27),
        21: ("left_wrist_yaw_joint", 29),
        26: ("right_wrist_roll_joint", 26),
        27: ("right_wrist_pitch_joint", 28),
        28: ("right_wrist_yaw_joint", 30),
    }
    for hardware_id, (joint_name, policy_id) in critical_joint_slots.items():
        if HARDWARE_JOINT_NAMES[hardware_id] != joint_name:
            raise RuntimeError(
                f"H2 joint-map name tripwire failed at hardware id {hardware_id}: "
                f"{HARDWARE_JOINT_NAMES[hardware_id]!r} != {joint_name!r}"
            )
        if int(contract.hw_to_policy[hardware_id]) != policy_id:
            raise RuntimeError(
                f"H2 joint-map policy tripwire failed at hardware id {hardware_id}: "
                f"{int(contract.hw_to_policy[hardware_id])} != {policy_id}"
            )
    if not np.allclose(ready_hw[29:31], 0.0):
        raise RuntimeError("head lock was not applied to official H2 motor ids 29/30")

    reference_joint_pos = frozen.joint_pos.copy()
    reference_joint_vel = np.zeros(H2_DOF, dtype=np.float64)
    reference_pelvis_pos = frozen.pelvis_pos.copy()
    reference_pelvis_quat = frozen.pelvis_quat.copy()
    generation = None
    if not policy_only:
        if settings.generator_device == "cuda" and not torch.cuda.is_available():
            raise RuntimeError("full preflight requested CUDA but torch.cuda.is_available() is false")
        generator = ReferenceGenerator(
            spec,
            list(contract.policy_names),
            contract.default_policy,
            device=settings.generator_device,
            sampling_steps=settings.generator_sampling_steps,
        )
        reference_state = RealReferenceState(
            spec, generator, frozen, landing_xy=settings.landing_xy,
            latency_skip_s=settings.active_latency_skip_s,
        )
        try:
            # Mirror the real runner: warm CUDA kernels, then exercise the
            # asynchronous request/result/latency-skip path with no robot I/O.
            reference_state.warmup(settings.generator_warmup_steps)
            reference_state.update_robot_state(
                np.zeros(2), np.array([1.0, 0.0, 0.0, 0.0]), frozen.joint_pos
            )
            reference_state.update_external_command(
                1, 0.80, np.array([0.40, 0.0, 1.10]), np.array([-3.0, 0.0, 0.2]), base_y=0.0
            )
            deadline = time.monotonic() + 30.0
            while reference_state.command_id != 1 and time.monotonic() < deadline:
                reference_state.step()
                time.sleep(0.005)
            reference_state.raise_if_failed()
            if reference_state.command_id != 1 or reference_state.is_frozen:
                raise RuntimeError("asynchronous H2 CLoSD result did not activate before timeout")
            arrays = [
                reference_state.ref_joint_pos,
                reference_state.ref_joint_vel,
                reference_state.ref_pelvis_pos,
                reference_state.ref_pelvis_quat,
            ]
            if not all(np.isfinite(array).all() for array in arrays):
                raise RuntimeError("asynchronous H2 CLoSD reference contains NaN/Inf")
            phase_values = np.asarray(
                [
                    reference_state.last_motion_hit_time_s,
                    reference_state.last_remaining_command_s,
                ],
                dtype=np.float64,
            )
            if not np.isfinite(phase_values).all():
                raise RuntimeError("H2 CLoSD phase alignment diagnostics are non-finite")
            phase_frames = int(
                round(
                    (
                        reference_state.last_motion_hit_time_s
                        - reference_state.last_remaining_command_s
                    )
                    * generator.fps
                )
            )
            expected_skip = phase_frames if phase_frames >= 1 else 1
            expected_wait = 0 if phase_frames >= 1 else 1 - phase_frames
            if reference_state.last_skip_frames != expected_skip:
                raise RuntimeError(
                    "H2 CLoSD phase skip invariant failed: "
                    f"{reference_state.last_skip_frames} != {expected_skip}"
                )
            if reference_state.last_wait_frames != expected_wait:
                raise RuntimeError(
                    "H2 CLoSD phase wait invariant failed: "
                    f"{reference_state.last_wait_frames} != {expected_wait}"
                )
            if reference_state.last_skip_frames < 1:
                raise RuntimeError("H2 CLoSD must always discard generated frame zero")
            if not np.isclose(
                reference_state.last_motion_hit_time_s,
                spec.condition_time_range[1],
                atol=1.0e-5,
            ):
                raise RuntimeError(
                    "the 0.80s preflight command was not clamped to the H2 bundle's "
                    f"effective hit time {spec.condition_time_range[1]:.3f}s"
                )
            reference_joint_pos = reference_state.ref_joint_pos.copy()
            reference_joint_vel = reference_state.ref_joint_vel.copy()
            reference_pelvis_pos = reference_state.ref_pelvis_pos.copy()
            reference_pelvis_quat = reference_state.ref_pelvis_quat.copy()
            generation = {
                "frames": generator.pred_len,
                "clamped": reference_state.last_clamped,
                "clamp_delta_m": reference_state.last_clamp_delta_m,
                "device": generator.device,
                "generation_s": reference_state.last_generation_s,
                "latency_skip_frames": reference_state.last_skip_frames,
                "skip_frames": reference_state.last_skip_frames,
                "wait_frames": reference_state.last_wait_frames,
                "motion_hit_s": reference_state.last_motion_hit_time_s,
                "remaining_command_s": reference_state.last_remaining_command_s,
                "async_path": True,
            }
        finally:
            reference_state.close()

    obs = build_observation(
        task_command=np.asarray(spec.frozen_command(), dtype=np.float64),
        ref_joint_pos=reference_joint_pos,
        ref_joint_vel=reference_joint_vel,
        root_xy=np.zeros(2),
        ref_pelvis_xy=reference_pelvis_pos[:2],
        ref_pelvis_quat=reference_pelvis_quat,
        pelvis_quat=np.array([1.0, 0.0, 0.0, 0.0]),
        heading_quat=np.array([1.0, 0.0, 0.0, 0.0]),
        base_ang_vel_body=np.zeros(3),
        joint_pos=frozen.joint_pos,
        joint_vel=np.zeros(H2_DOF),
        default_joint_pos=contract.default_policy,
        prev_action=np.zeros(H2_DOF),
        noise=False,
    )
    if obs.shape != (expected_obs_dim(H2_DOF),) or not np.isfinite(obs).all():
        raise RuntimeError(f"invalid H2 observation: {obs.shape}")
    action = np.clip(policy(obs), -10.0, 10.0)
    # Exercise the actual learned-policy contract: the clipped action maps to
    # a raw virtual PD target.  Mechanical position limiting is reserved for
    # ready/stop; the bridge performs the latest-state effort projection.
    target_hw = contract.learned_raw_hw_targets_from_action(action)
    if not np.isfinite(target_hw).all() or not np.allclose(target_hw[29:31], 0.0):
        raise RuntimeError("invalid policy-to-H2 target conversion")

    return {
        "deployment_version": settings.deployment_version,
        "policy_path": str(settings.policy_path),
        "status": "PASS",
        "artifacts_verified": len(verified),
        "policy_obs_dim": policy.obs_dim,
        "policy_action_dim": policy.act_dim,
        "shm_bytes": BRIDGE_SIZE,
        "shm_version": SHM_VERSION,
        "mode_pr": 0,
        "joint_map_revision": H2_JOINT_MAP_REVISION,
        "hw_to_policy": contract.hw_to_policy.tolist(),
        "policy_to_hw": contract.policy_to_hw.tolist(),
        "critical_joint_slots_by_hw_id": {
            str(hardware_id): {
                "name": joint_name,
                "policy_index": policy_id,
            }
            for hardware_id, (joint_name, policy_id) in critical_joint_slots.items()
        },
        "head_target_hw_29_30": target_hw[29:31].tolist(),
        "generation": generation,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=ROOT / "config.yaml")
    parser.add_argument("--deployment-version", "--upstream-version", choices=DEPLOYMENTS)
    parser.add_argument("--policy-only", action="store_true", help="Skip the native CLoSD generation smoke test.")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()
    report = run(args.config, policy_only=args.policy_only, deployment_version=args.deployment_version)
    if args.json:
        print(json.dumps(report, indent=2, sort_keys=True))
    else:
        print("[PREFLIGHT PASS] " + json.dumps(report, sort_keys=True))


if __name__ == "__main__":
    main()