"""`Okvis2Strategy` — production-default C1 VIO (AZ-332).

Python facade over the OKVIS2 C++ tightly-coupled keyframe-based VIO
core, accessed via the pybind11 binding at
``_native.okvis2_binding.Okvis2Backend`` (compiled by
``cpp/okvis2/CMakeLists.txt``, gated by ``BUILD_OKVIS2=ON``).

Conforms to the AZ-331 :class:`VioStrategy` Protocol; consumes the
runtime ``Config`` + an :class:`FdrClient`; constructs its other
dependencies (logger, camera calibration, preintegrator) internally
from ``config`` so the strategy class matches the composition-root
factory shape::

    strategy_cls(config: Config, *, fdr_client: FdrClient)

Risk-2 mitigation: the native binding is imported **lazily inside the
constructor**, not at module top level. Importing this module with
``BUILD_OKVIS2=OFF`` (no compiled ``.so``) is safe — the factory's
build-flag gate catches that path before the constructor runs.

AC mapping (see ``_docs/02_tasks/todo/AZ-332_c1_okvis2_strategy.md``):

- AC-1 : :meth:`current_strategy_label` returns ``"okvis2"``.
- AC-2 : :meth:`process_frame` returns :class:`VioOutput` with
  ``frame_id`` echoed; covariance SPD; ``imu_bias`` non-None.
- AC-3 : all backend / Eigen / std::runtime_error rewrap into
  :class:`VioError` family with ``__cause__`` chain.
- AC-4 : :meth:`reset_to_warm_start` clears state + seeds hint; second
  consecutive call does not raise.
- AC-5 : :meth:`health_snapshot` returns INIT initially, TRACKING after
  ``warm_start_max_frames`` (default 5) successful frames.
- AC-6 : DEGRADED on feature loss; covariance Frobenius norm strictly
  increases; ``process_frame`` still returns :class:`VioOutput` (not raise).
- AC-7 : after ``lost_frame_threshold`` (default 9) consecutive failed
  frames, raises :class:`VioFatalError`; state == LOST.
- AC-8 : ``BUILD_OKVIS2=OFF`` does not load this module (enforced by
  AZ-331 factory; covered in
  ``tests/unit/c1_vio/test_protocol_conformance.py``).
- AC-9 / NFR-perf : tier2 — Jetson + Derkachi-class fixture; tests
  marked ``@pytest.mark.tier2``.
- AC-10 : exactly one ``vio.health`` FDR record per state transition;
  no spam on steady-state.
"""

from __future__ import annotations

import math
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final, Literal

import numpy as np

from gps_denied_onboard._types.nav import (
    FeatureQuality,
    ImuBias,
    VioHealth,
    VioOutput,
    VioState,
)
from gps_denied_onboard.clock.wall_clock import WallClock
from gps_denied_onboard.components.c1_vio.errors import (
    VioFatalError,
    VioInitializingError,
)
from gps_denied_onboard.fdr_client.records import CURRENT_SCHEMA_VERSION, FdrRecord
from gps_denied_onboard.logging import get_logger

if TYPE_CHECKING:
    import numpy.typing as npt

    from gps_denied_onboard._types.calibration import CameraCalibration
    from gps_denied_onboard._types.nav import (
        ImuWindow,
        NavCameraFrame,
        WarmStartPose,
    )
    from gps_denied_onboard.clock import Clock
    from gps_denied_onboard.components.c1_vio.config import Okvis2Config
    from gps_denied_onboard.config import Config
    from gps_denied_onboard.fdr_client.client import FdrClient

__all__ = ["Okvis2Strategy"]


_STRATEGY_LABEL: Final[Literal["okvis2"]] = "okvis2"
_PRODUCER_ID: Final[str] = "c1_vio.okvis2"
_LOGGER_COMPONENT: Final[str] = "c1_vio.okvis2"
_BIAS_NORM_FLOOR: Final[float] = 0.0


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _bias_norm(bias: ImuBias) -> float:
    """L2 norm of the concatenated 6-vector ``(accel || gyro)``."""
    accel = np.asarray(bias.accel_bias, dtype=np.float64)
    gyro = np.asarray(bias.gyro_bias, dtype=np.float64)
    return float(math.sqrt(float(np.dot(accel, accel) + np.dot(gyro, gyro))))


def _se3_from_4x4(matrix: npt.NDArray[Any]) -> Any:
    """Build a ``gtsam.Pose3`` from a 4x4 row-major matrix.

    Imported lazily so this module can be imported without gtsam in
    headless tooling paths (tests + facade-only smoke).
    """
    import gtsam

    return gtsam.Pose3(np.asarray(matrix, dtype=np.float64))


class Okvis2Strategy:
    """Production-default :class:`VioStrategy` for E-C1 (AZ-332).

    Constructor matches the AZ-331 composition-root factory shape::

        Okvis2Strategy(config: Config, *, fdr_client: FdrClient)

    Other dependencies (calibration, preintegrator-substrate, logger,
    OKVIS2 sub-config) are resolved internally from ``config``.

    Concurrency: single-threaded by Protocol invariant. One instance
    per camera-ingest writer thread; concurrent ``process_frame`` calls
    are undefined behaviour.
    """

    def __init__(
        self,
        config: Config,
        *,
        fdr_client: FdrClient,
        clock: Clock | None = None,
    ) -> None:
        c1_block = config.components["c1_vio"]
        if c1_block.strategy != _STRATEGY_LABEL:
            raise VioFatalError(
                f"Okvis2Strategy constructed with config.strategy="
                f"{c1_block.strategy!r}; expected {_STRATEGY_LABEL!r}. "
                "The AZ-331 factory is the only sanctioned constructor."
            )

        self._config = config
        self._fdr = fdr_client
        self._clock: Clock = clock if clock is not None else WallClock()
        self._logger = get_logger(_LOGGER_COMPONENT)
        self._lost_frame_threshold: int = c1_block.lost_frame_threshold
        self._warm_start_max_frames: int = c1_block.warm_start_max_frames
        self._okvis2_cfg: Okvis2Config = c1_block.okvis2
        self._calibration: CameraCalibration | None = None
        self._frames_since_warmup: int = 0
        self._consecutive_lost: int = 0
        self._latest_bias: ImuBias = ImuBias(accel_bias=(0.0, 0.0, 0.0), gyro_bias=(0.0, 0.0, 0.0))
        self._reported_state: VioState = VioState.INIT
        self._last_emitted_state: VioState | None = None

        # Lazy import of the native binding — Risk-2 mitigation (I-5).
        # Failure here is the BUILD_OKVIS2=OFF path the AZ-331 factory's
        # `StrategyNotAvailableError` is meant to prevent; if a caller
        # bypasses the factory and reaches this constructor with the
        # native lib absent, we surface a fatal init error.
        try:
            from gps_denied_onboard.components.c1_vio._native import (
                okvis2_binding,
            )
        except ImportError as exc:
            raise VioFatalError(
                "Okvis2Strategy: native binding "
                "(gps_denied_onboard.components.c1_vio._native.okvis2_binding) "
                "is not importable. Production binary must be built with "
                "BUILD_OKVIS2=ON."
            ) from exc

        self._binding_module = okvis2_binding
        self._backend = self._construct_backend()

    # ------------------------------------------------------------------
    # Public Protocol surface.

    def process_frame(
        self,
        frame: NavCameraFrame,
        imu: ImuWindow,
        calibration: CameraCalibration,
    ) -> VioOutput:
        """Hot-path call — one per nav-camera frame.

        Steps:

        1. Push every IMU sample in the window into the backend; the
           strict-monotonic guard lives on the C++ side.
        2. Submit the frame.
        3. If the backend produced an output, classify health and
           build the :class:`VioOutput` DTO.
        4. If no output: tick the lost-frame counter; emit a state
           transition record if needed.
        """
        self._calibration = calibration
        frame_id_str = str(frame.frame_id)
        emitted_at_ns = self._clock.monotonic_ns()

        try:
            self._push_imu_window(imu)
            produced = self._backend.add_frame(
                frame_id_str, _frame_ts_ns(frame), _frame_image(frame)
            )
        except self._binding_module.OkvisInitException as exc:
            self._emit_transition(VioState.INIT, frame_id_str)
            raise VioInitializingError(
                f"OKVIS2 backend reports INIT while processing frame {frame_id_str!r}: {exc}"
            ) from exc
        except self._binding_module.OkvisOptimizationException as exc:
            # Treat as a degraded frame: emit no VioOutput from this
            # path — callers expect either a VioOutput or a VioError;
            # we choose error here so C5 can fall back, matching AC-3.
            self._tick_lost(frame_id_str)
            if self._reported_state == VioState.LOST:
                self._emit_transition(VioState.LOST, frame_id_str)
                raise VioFatalError(
                    f"OKVIS2 backend exhausted lost-frame budget at {frame_id_str!r}: {exc}"
                ) from exc
            self._emit_transition(self._reported_state, frame_id_str)
            raise VioInitializingError(
                f"OKVIS2 backend optimisation failure at {frame_id_str!r}: {exc}"
            ) from exc
        except self._binding_module.OkvisFatalException as exc:
            self._emit_transition(VioState.LOST, frame_id_str)
            raise VioFatalError(
                f"OKVIS2 backend fatal exception at {frame_id_str!r}: {exc}"
            ) from exc
        except (RuntimeError, ValueError) as exc:
            # Catch-all for unmapped backend exceptions. Re-classify as
            # fatal — we explicitly forbid raw library exceptions across
            # the public boundary.
            raise VioFatalError(
                f"OKVIS2 backend raised an unmapped exception at {frame_id_str!r}: {exc}"
            ) from exc

        if not produced:
            # Frame consumed but no estimator update yet — INIT path
            # while OKVIS2 warms up its keyframe window.
            self._emit_transition(VioState.INIT, frame_id_str)
            raise VioInitializingError(
                f"Okvis2Strategy: backend has not yet emitted an "
                f"estimator update at {frame_id_str!r}"
            )

        raw = self._backend.get_latest_output()
        if raw is None:
            raise VioFatalError(
                f"Okvis2Strategy: backend reported a new output for "
                f"{frame_id_str!r} but get_latest_output() returned None"
            )

        vio_output = self._build_vio_output(raw, emitted_at_ns)
        self._consecutive_lost = 0
        new_state = self._classify_state(vio_output.feature_quality)
        if new_state != self._reported_state:
            self._reported_state = new_state
            self._emit_transition(new_state, frame_id_str)

        if new_state in (VioState.INIT, VioState.TRACKING):
            self._frames_since_warmup += 1

        if self._okvis2_cfg.per_frame_debug_log:
            self._logger.debug(
                "okvis2.process_frame",
                extra={
                    "component": _LOGGER_COMPONENT,
                    "kind": "vio.tick",
                    "frame_id": frame_id_str,
                    "kv": {
                        "state": self._reported_state.value,
                        "tracked": vio_output.feature_quality.tracked,
                        "mre_px": vio_output.feature_quality.mre_px,
                        "emitted_at_ns": vio_output.emitted_at_ns,
                    },
                },
            )

        return vio_output

    def reset_to_warm_start(self, hint: WarmStartPose) -> None:
        """Destructive re-init from an F8-reboot warm-start hint.

        Idempotent across consecutive calls (AC-4) — a second call
        without an intervening ``process_frame`` reseats the backend
        again without raising.
        """
        try:
            body_T_world = np.asarray(hint.body_T_world.matrix(), dtype=np.float64)
        except AttributeError as exc:
            raise VioFatalError(
                "Okvis2Strategy.reset_to_warm_start: hint.body_T_world is "
                "not a gtsam.Pose3 (missing .matrix())"
            ) from exc

        velocity = np.asarray(hint.velocity_b, dtype=np.float64)
        accel_bias = np.asarray(hint.bias.accel_bias, dtype=np.float64)
        gyro_bias = np.asarray(hint.bias.gyro_bias, dtype=np.float64)

        try:
            self._backend.reset(body_T_world, velocity, accel_bias, gyro_bias)
        except self._binding_module.OkvisInitException as exc:
            raise VioFatalError(f"OKVIS2 backend rejected warm-start reset: {exc}") from exc
        except (RuntimeError, ValueError) as exc:
            raise VioFatalError(
                f"OKVIS2 backend raised an unmapped exception during reset: {exc}"
            ) from exc

        self._latest_bias = hint.bias
        self._frames_since_warmup = 0
        self._consecutive_lost = 0
        self._reported_state = VioState.INIT
        self._emit_transition(VioState.INIT, frame_id="")

    def health_snapshot(self) -> VioHealth:
        """Most-recent health state — no backend call (cheap)."""
        return VioHealth(
            state=self._reported_state,
            consecutive_lost=self._consecutive_lost,
            bias_norm=_bias_norm(self._latest_bias),
        )

    def current_strategy_label(self) -> Literal["okvis2", "vins_mono", "klt_ransac"]:
        return _STRATEGY_LABEL

    # ------------------------------------------------------------------
    # Internal helpers.

    def _construct_backend(self) -> Any:
        """Build the backend from config — calibration path is optional
        because the unit-test fake-binding path skips real intrinsics.

        Tests inject a fake module at ``sys.modules`` before construction
        (see ``tests/unit/c1_vio/conftest.py``); the fake's
        ``Okvis2Backend`` accepts whatever this method passes.
        """
        K = self._load_camera_intrinsics()
        yaml_config = self._render_yaml_config()
        try:
            return self._binding_module.Okvis2Backend(yaml_config, K)
        except self._binding_module.OkvisInitException as exc:
            raise VioFatalError(f"Okvis2Strategy: backend init failed: {exc}") from exc

    def _load_camera_intrinsics(self) -> np.ndarray:
        """Load 3x3 camera intrinsics from the calibration path.

        Returns the identity matrix when the runtime block has no
        path configured — the unit-test path overrides this via the
        fake binding's ctor anyway, and a production binary refusing
        to start on a missing calibration is preferable to silently
        emitting wrong poses (handled by the YAML loader downstream).
        """
        path = self._config.runtime.camera_calibration_path
        if not path:
            return np.eye(3, dtype=np.float64)
        try:
            import json

            with open(path, encoding="utf-8") as fh:
                blob = json.load(fh)
        except (OSError, ValueError) as exc:
            raise VioFatalError(
                f"Okvis2Strategy: failed to load camera calibration from {path!r}: {exc}"
            ) from exc
        K_raw = blob.get("intrinsics_3x3")
        if K_raw is None:
            raise VioFatalError(
                f"Okvis2Strategy: calibration file {path!r} is missing the 'intrinsics_3x3' field"
            )
        K = np.asarray(K_raw, dtype=np.float64)
        if K.shape != (3, 3):
            raise VioFatalError(f"Okvis2Strategy: intrinsics_3x3 must be 3x3; got shape {K.shape}")
        return K

    def _render_yaml_config(self) -> str:
        """Render the Okvis2Config sub-block into an OKVIS2 YAML snippet.

        OKVIS2 reads a YAML config string at construction. Only the knobs
        AZ-332 exposes are rendered; OKVIS2-internal defaults cover the
        rest.
        """
        cfg = self._okvis2_cfg
        return (
            "# AZ-332 — generated OKVIS2 config (see Okvis2Config in c1_vio/config.py)\n"
            f"keyframe_window_size: {cfg.keyframe_window_size}\n"
            f"keyframe_parallax_threshold_px: {cfg.keyframe_parallax_threshold_px}\n"
            f"ransac_inlier_ratio: {cfg.ransac_inlier_ratio}\n"
            f"max_optimization_iters: {cfg.max_optimization_iters}\n"
        )

    def _push_imu_window(self, imu: ImuWindow) -> None:
        for sample in imu.samples:
            self._backend.add_imu(
                sample.ts_ns,
                np.asarray(sample.accel_xyz, dtype=np.float64),
                np.asarray(sample.gyro_xyz, dtype=np.float64),
            )

    def _build_vio_output(self, raw: dict[str, Any], emitted_at_ns: int) -> VioOutput:
        try:
            pose = _se3_from_4x4(raw["pose_T_world_body"])
            cov = np.asarray(raw["pose_covariance_6x6"], dtype=np.float64)
            bias = ImuBias(
                accel_bias=tuple(float(x) for x in raw["accel_bias"]),  # type: ignore[arg-type]
                gyro_bias=tuple(float(x) for x in raw["gyro_bias"]),  # type: ignore[arg-type]
            )
            feature_quality = FeatureQuality(
                tracked=int(raw["tracked_features"]),
                new=int(raw["new_features"]),
                lost=int(raw["lost_features"]),
                mean_parallax=float(raw["mean_parallax"]),
                mre_px=float(raw["mre_px"]),
            )
            backend_ts = int(raw.get("emitted_at_ns") or emitted_at_ns)
        except (KeyError, TypeError, ValueError) as exc:
            raise VioFatalError(f"Okvis2Strategy: backend output is malformed: {exc}") from exc

        if cov.shape != (6, 6):
            raise VioFatalError(
                f"Okvis2Strategy: pose_covariance_6x6 has shape {cov.shape}; expected (6, 6)"
            )

        self._latest_bias = bias
        return VioOutput(
            frame_id=raw["frame_id"],
            relative_pose_T=pose,
            pose_covariance_6x6=cov,
            imu_bias=bias,
            feature_quality=feature_quality,
            emitted_at_ns=backend_ts,
        )

    def _classify_state(self, fq: FeatureQuality) -> VioState:
        if self._reported_state == VioState.INIT and (
            self._frames_since_warmup + 1 < self._warm_start_max_frames
        ):
            return VioState.INIT
        if fq.tracked < self._okvis2_cfg.degraded_feature_threshold:
            return VioState.DEGRADED
        return VioState.TRACKING

    def _tick_lost(self, frame_id: str) -> None:
        self._consecutive_lost += 1
        if self._consecutive_lost >= self._lost_frame_threshold:
            self._reported_state = VioState.LOST
        elif self._reported_state == VioState.TRACKING:
            self._reported_state = VioState.DEGRADED

    def _emit_transition(self, new_state: VioState, frame_id: str) -> None:
        if self._last_emitted_state == new_state:
            return
        self._last_emitted_state = new_state
        record = FdrRecord(
            schema_version=CURRENT_SCHEMA_VERSION,
            ts=_now_iso(),
            producer_id=_PRODUCER_ID,
            kind="vio.health",
            payload={
                "state": new_state.value,
                "consecutive_lost": self._consecutive_lost,
                "bias_norm": _bias_norm(self._latest_bias),
                "strategy_label": _STRATEGY_LABEL,
                "frame_id": frame_id,
            },
        )
        self._fdr.enqueue(record)


def _frame_ts_ns(frame: NavCameraFrame) -> int:
    """Convert ``NavCameraFrame.timestamp`` to monotonic-ns.

    Uses the datetime's UTC epoch nanoseconds so the value is
    monotonically increasing across frames (frame source guarantees
    strictly increasing timestamps per the FrameSource contract).
    """
    return int(frame.timestamp.timestamp() * 1e9)


def _frame_image(frame: NavCameraFrame) -> np.ndarray:
    """Coerce the frame's image into a contiguous uint8 ndarray."""
    arr = np.ascontiguousarray(frame.image, dtype=np.uint8)
    if arr.ndim < 2 or arr.ndim > 3:
        raise VioFatalError(
            f"Okvis2Strategy: NavCameraFrame.image must be 2-D or 3-D; got {arr.ndim}-D"
        )
    return arr
