"""``python -m gps_denied_onboard.components.c1_vio.bench.okvis2`` (AZ-332).

Microbench for :class:`Okvis2Strategy` — reads a fixture directory of
nav-camera frames + IMU samples and reports per-frame latency
percentiles for C1-PT-01 (p50 <= 25 ms, p95 <= 80 ms, threshold 120 ms).

The bench produces production behaviour: it constructs the real
strategy via the AZ-331 factory (so ``BUILD_OKVIS2=ON`` is required),
feeds real frames through, and measures wall-clock per call. On Tier-2
this measures OKVIS2's actual estimator latency; on a workstation with
``BUILD_OKVIS2=OFF`` it refuses to start (Risk-2 — never silently
benchmark a stub).
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np

from gps_denied_onboard._types.nav import (
    ImuSample,
    ImuWindow,
    NavCameraFrame,
)
from gps_denied_onboard.components.c1_vio.config import (
    C1VioConfig,
    Okvis2Config,
)
from gps_denied_onboard.config.schema import Config, RuntimeConfig
from gps_denied_onboard.fdr_client.client import make_fdr_client
from gps_denied_onboard.runtime_root.vio_factory import build_vio_strategy


def _percentile(samples_ms: list[float], pct: float) -> float:
    if not samples_ms:
        return float("nan")
    sorted_samples = sorted(samples_ms)
    idx = min(len(sorted_samples) - 1, int(pct * len(sorted_samples)))
    return sorted_samples[idx]


def _load_fixture(fixture_dir: Path) -> tuple[Any, list[NavCameraFrame], list[ImuWindow]]:
    """Fixture format (minimal, deterministic):

    .. code::

        fixture_dir/
          manifest.json    { "frame_count": N, "camera_calibration_path": "..." }
          frames/0000.npy  uint8 image
          ...
          imu/0000.json    {"samples": [{"ts_ns": N, "accel": [..], "gyro": [..]}, ...]}
          ...
    """
    manifest_path = fixture_dir / "manifest.json"
    if not manifest_path.is_file():
        raise FileNotFoundError(f"missing manifest.json under {fixture_dir!r}")
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))

    frames: list[NavCameraFrame] = []
    imu_windows: list[ImuWindow] = []
    frame_count = int(manifest["frame_count"])
    for i in range(frame_count):
        img_path = fixture_dir / "frames" / f"{i:04d}.npy"
        imu_path = fixture_dir / "imu" / f"{i:04d}.json"
        img = np.load(img_path)
        imu_blob = json.loads(imu_path.read_text(encoding="utf-8"))
        samples = tuple(
            ImuSample(
                ts_ns=int(s["ts_ns"]),
                accel_xyz=tuple(s["accel"]),
                gyro_xyz=tuple(s["gyro"]),
            )
            for s in imu_blob["samples"]
        )
        if not samples:
            raise ValueError(
                f"bench.okvis2: fixture frame {i} ({imu_path}) has no IMU "
                "samples — bench requires a real IMU window per frame"
            )
        ts_start = samples[0].ts_ns
        ts_end = samples[-1].ts_ns
        imu_windows.append(ImuWindow(samples=samples, ts_start_ns=ts_start, ts_end_ns=ts_end))
        frames.append(
            NavCameraFrame(
                frame_id=i,
                timestamp=datetime.fromtimestamp(ts_start * 1e-9, tz=timezone.utc),
                image=img,
                camera_calibration_id=str(manifest.get("camera_calibration_id", "bench")),
            )
        )
    return manifest, frames, imu_windows


def _make_calibration(intrinsics_path: str | None) -> Any:
    """Build a CameraCalibration with no body-to-camera (identity)
    using the bench's calibration JSON if supplied; otherwise raise.
    """
    from gps_denied_onboard._types.calibration import CameraCalibration

    if intrinsics_path is None:
        raise ValueError("bench.okvis2: --camera-calibration is required (real intrinsics)")
    blob = json.loads(Path(intrinsics_path).read_text(encoding="utf-8"))
    return CameraCalibration(
        camera_id=blob.get("camera_id", "bench"),
        intrinsics_3x3=np.asarray(blob["intrinsics_3x3"], dtype=np.float64),
        distortion=np.asarray(blob.get("distortion", [0, 0, 0, 0]), dtype=np.float64),
        body_to_camera_se3=np.eye(4, dtype=np.float64),
        acquisition_method=blob.get("acquisition_method", "bench-static"),
        metadata=dict(blob.get("metadata", {})),
    )


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        prog="python -m gps_denied_onboard.components.c1_vio.bench.okvis2",
        description="Microbench for Okvis2Strategy.process_frame (AZ-332 / C1-PT-01).",
    )
    parser.add_argument("fixture_dir", type=Path, help="Path to fixture directory")
    parser.add_argument(
        "--camera-calibration",
        type=str,
        required=True,
        help="Path to camera calibration JSON",
    )
    parser.add_argument(
        "--warmup",
        type=int,
        default=10,
        help="Number of warmup frames (not counted in percentiles)",
    )
    args = parser.parse_args(argv)

    manifest, frames, imu_windows = _load_fixture(args.fixture_dir)
    calibration = _make_calibration(args.camera_calibration)

    config = Config.with_blocks(
        c1_vio=C1VioConfig(strategy="okvis2", okvis2=Okvis2Config()),
        runtime=RuntimeConfig(
            camera_calibration_path=args.camera_calibration,
            inference_backend="tensorrt",
            tier=2,
        ),
    )
    fdr_client = make_fdr_client("c1_vio.okvis2.bench", config)
    strategy = build_vio_strategy(config, fdr_client=fdr_client)

    durations_ms: list[float] = []
    for i, (frame, imu) in enumerate(zip(frames, imu_windows, strict=True)):
        t0 = time.perf_counter()
        try:
            strategy.process_frame(frame, imu, calibration)
        except Exception as exc:
            print(
                f"frame {i}: exception {type(exc).__name__}: {exc}",
                file=sys.stderr,
            )
            continue
        dt_ms = (time.perf_counter() - t0) * 1000.0
        if i >= args.warmup:
            durations_ms.append(dt_ms)

    if not durations_ms:
        print("bench: no successful frames after warmup", file=sys.stderr)
        return 2

    p50 = _percentile(durations_ms, 0.50)
    p95 = _percentile(durations_ms, 0.95)
    p99 = _percentile(durations_ms, 0.99)
    print(
        json.dumps(
            {
                "fixture_dir": str(args.fixture_dir),
                "frame_count": manifest.get("frame_count"),
                "measured": len(durations_ms),
                "p50_ms": round(p50, 3),
                "p95_ms": round(p95, 3),
                "p99_ms": round(p99, 3),
                "c1_pt_01_target_p50_ms": 25.0,
                "c1_pt_01_target_p95_ms": 80.0,
                "c1_pt_01_failure_p95_ms": 120.0,
            },
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
