"""E2E harness — runs a DatasetAdapter through FlightProcessor and collects output.

The harness instantiates `FlightProcessor` with a minimal in-memory repository
and an event sink instead of the real SSE streamer. It pushes adapter frames
through `process_frame()` and records each emitted FrameResult for comparison
against the adapter's ground truth.

What the harness does NOT do:
  - Train/tune ESKF noise parameters (use adapter-specific defaults)
  - SE(3) trajectory alignment (callers apply as needed)
  - Image decoding beyond OpenCV (extended formats handled per-adapter)
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
from unittest.mock import AsyncMock, MagicMock

import cv2
import numpy as np

from gps_denied.core.chunk_manager import RouteChunkManager
from gps_denied.core.coordinates import CoordinateTransformer
from gps_denied.core.gpr import GlobalPlaceRecognition
from gps_denied.core.graph import FactorGraphOptimizer
from gps_denied.core.metric import MetricRefinement
from gps_denied.core.models import ModelManager
from gps_denied.core.processor import FlightProcessor
from gps_denied.core.recovery import FailureRecoveryCoordinator
from gps_denied.core.vo import ISequentialVisualOdometry, ORBVisualOdometry
from gps_denied.schemas import GPSPoint
from gps_denied.schemas.graph import FactorGraphConfig
from gps_denied.schemas.vo import RelativePose
from gps_denied.testing.datasets.base import (
    DatasetAdapter,
    DatasetPose,
    PlatformClass,
)

EARTH_R = 6_378_137.0


class _ScaledVO(ISequentialVisualOdometry):
    """Thin wrapper that multiplies the unit-scale ORB translation by a fixed factor.

    ORB `recoverPose` returns a unit vector for translation — metric scale is
    lost in the Essential Matrix decomposition.  When a rough per-frame distance
    is known (e.g. from IMU or dataset spec) this wrapper rescales so the ESKF
    receives translations in metres instead of arbitrary units.
    """

    def __init__(self, inner: ISequentialVisualOdometry, scale_m: float) -> None:
        self._inner = inner
        self._scale_m = scale_m

    def extract_features(self, image):
        return self._inner.extract_features(image)

    def match_features(self, features1, features2):
        return self._inner.match_features(features1, features2)

    def estimate_motion(self, matches, camera_params):
        return self._inner.estimate_motion(matches, camera_params)

    def compute_relative_pose(self, prev_image, curr_image, camera_params):
        pose = self._inner.compute_relative_pose(prev_image, curr_image, camera_params)
        if pose is None:
            return None
        norm = float(np.linalg.norm(pose.translation))
        scaled_t = pose.translation / max(norm, 1e-9) * self._scale_m
        return RelativePose(
            translation=scaled_t,
            rotation=pose.rotation,
            confidence=pose.confidence,
            inlier_count=pose.inlier_count,
            total_matches=pose.total_matches,
            tracking_good=pose.tracking_good,
            scale_ambiguous=False,
        )


@dataclass
class HarnessResult:
    num_frames_submitted: int
    num_estimates: int
    estimated_positions_enu: np.ndarray = field(default_factory=lambda: np.zeros((0, 3)))
    ground_truth: np.ndarray = field(default_factory=lambda: np.zeros((0, 3)))
    eskf_positions_enu: np.ndarray = field(default_factory=lambda: np.zeros((0, 3)))
    adapter_name: str = ""
    platform_class: PlatformClass = PlatformClass.SYNTHETIC


class E2EHarness:
    """Drives FlightProcessor from a DatasetAdapter; collects results.

    `max_frames` caps the iteration (and the matching GT slice). Useful for:
      - CI tiers where a full multi-thousand-frame sequence is too slow.
      - Debugging a new adapter — eyeball the first N frames before committing
        to a multi-hour run.
    `None` (default) means "consume the full dataset".
    """

    def __init__(
        self,
        adapter: DatasetAdapter,
        flight_id: str = "e2e-flight",
        max_frames: Optional[int] = None,
        trace_path: Optional[Path] = None,
        vo_scale_m: Optional[float] = None,
    ) -> None:
        self._adapter = adapter
        self._flight_id = flight_id
        self._max_frames = max_frames
        self._trace_path: Optional[Path] = Path(trace_path) if trace_path else None
        self._vo_scale_m = vo_scale_m
        self._estimates: list[tuple[int, Optional[tuple[float, float, float]]]] = []

    async def run(self) -> HarnessResult:
        processor = self._build_processor()
        frames = list(self._adapter.iter_frames())
        gt_poses = list(self._adapter.iter_ground_truth())

        if self._max_frames is not None:
            frames = frames[: self._max_frames]
            gt_poses = gt_poses[: self._max_frames]

        # Align GT by index so trace records can pair each frame with the
        # corresponding pose without an expensive timestamp search.
        gt_by_idx: dict[int, DatasetPose] = {}
        for i, pose in enumerate(gt_poses):
            gt_by_idx[i] = pose

        # Seed ESKF with first GT pose as synthetic GPS origin so the processor
        # can propagate VO translations into a consistent ENU frame.
        if gt_poses:
            origin = gt_poses[0]
            start_gps = GPSPoint(lat=origin.lat, lon=origin.lon)
            processor._coord.set_enu_origin(self._flight_id, start_gps)  # noqa: SLF001
            processor._init_eskf_for_flight(  # noqa: SLF001
                self._flight_id, start_gps, altitude=origin.alt
            )

        trace_fh = None
        if self._trace_path is not None:
            self._trace_path.parent.mkdir(parents=True, exist_ok=True)
            trace_fh = self._trace_path.open("w")

        eskf_enu_rows: list[list[float]] = []

        try:
            for frame in frames:
                image = self._load_or_synth_image(frame.image_path)
                result = await processor.process_frame(
                    self._flight_id, frame.frame_idx, image
                )
                est = None
                if result.gps is not None:
                    est = (result.gps.lat, result.gps.lon, 0.0)
                self._estimates.append((frame.frame_idx, est))

                # Collect raw ESKF ENU position for drift analysis.
                eskf = processor._eskf.get(self._flight_id)  # noqa: SLF001
                if eskf and eskf.initialized:
                    pos = eskf.position
                    eskf_enu_rows.append([float(pos[0]), float(pos[1]), float(pos[2])])

                if trace_fh is not None:
                    gt = gt_by_idx.get(frame.frame_idx)
                    record = self._trace_record(processor, frame, result, gt)
                    trace_fh.write(json.dumps(record) + "\n")
        finally:
            if trace_fh is not None:
                trace_fh.close()

        gt_enu = self._poses_to_enu(gt_poses)
        est_enu = self._estimates_to_enu(gt_poses[0] if gt_poses else None)
        eskf_enu = np.array(eskf_enu_rows) if eskf_enu_rows else np.zeros((0, 3))

        return HarnessResult(
            num_frames_submitted=len(frames),
            num_estimates=sum(1 for _, e in self._estimates if e is not None),
            estimated_positions_enu=est_enu,
            ground_truth=gt_enu,
            eskf_positions_enu=eskf_enu,
            adapter_name=self._adapter.name,
            platform_class=self._adapter.capabilities.platform_class,
        )

    def _trace_record(
        self,
        processor: FlightProcessor,
        frame,
        result,
        gt: Optional[DatasetPose],
    ) -> dict:
        """Build one JSONL record describing the product's state after a frame.

        Captures VO success, ESKF state (position + trace(cov) as a scalar
        sigma), and the estimate/GT pair. Enough to diagnose *where* the
        pipeline diverges without dumping raw images or covariance matrices.
        """
        eskf = processor._eskf.get(self._flight_id)  # noqa: SLF001 — test harness
        eskf_initialized = bool(eskf and eskf.initialized)
        eskf_position = None
        eskf_sigma = None
        if eskf_initialized:
            pos = eskf.position
            eskf_position = [float(pos[0]), float(pos[1]), float(pos[2])]
            cov = eskf.covariance
            if cov is not None:
                # Scalar horizontal+vertical position uncertainty summary.
                eskf_sigma = float(np.sqrt(np.trace(cov[0:3, 0:3]) / 3.0))

        est_lat = result.gps.lat if result.gps is not None else None
        est_lon = result.gps.lon if result.gps is not None else None

        return {
            "frame_idx": frame.frame_idx,
            "timestamp_ns": frame.timestamp_ns,
            "vo_success": bool(result.vo_success),
            "alignment_success": bool(result.alignment_success),
            "tracking_state": result.tracking_state.value,
            "confidence": float(result.confidence),
            "eskf_initialized": eskf_initialized,
            "eskf_position_enu": eskf_position,
            "eskf_pos_sigma_m": eskf_sigma,
            "estimate_lat": est_lat,
            "estimate_lon": est_lon,
            "gt_lat": gt.lat if gt else None,
            "gt_lon": gt.lon if gt else None,
            "gt_alt": gt.alt if gt else None,
        }

    def _build_processor(self) -> FlightProcessor:
        repo = MagicMock()
        streamer = MagicMock()
        streamer.push_event = AsyncMock()
        proc = FlightProcessor(repo, streamer)
        mm = ModelManager()
        vo = ORBVisualOdometry()
        gpr = GlobalPlaceRecognition(mm)
        gpr.load_index(self._flight_id, "dummy")
        metric = MetricRefinement(mm)
        graph = FactorGraphOptimizer(FactorGraphConfig())
        chunk_mgr = RouteChunkManager(graph)
        recovery = FailureRecoveryCoordinator(chunk_mgr, gpr, metric)
        if self._vo_scale_m is not None:
            vo = _ScaledVO(vo, self._vo_scale_m)
        coord = CoordinateTransformer()
        proc.attach_components(
            vo=vo, gpr=gpr, metric=metric,
            graph=graph, recovery=recovery, chunk_mgr=chunk_mgr,
            coord=coord,
        )
        return proc

    def _load_or_synth_image(self, path: str) -> np.ndarray:
        if path:
            img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
            if img is None:
                raise FileNotFoundError(f"Could not read {path}")
            return img
        # Synthetic adapter: generate a checkerboard so OpenCV / VO have corners
        sz = 480
        board = np.zeros((sz, sz), dtype=np.uint8)
        cell = 40
        for r in range(0, sz, cell):
            for c in range(0, sz, cell):
                if ((r // cell) + (c // cell)) % 2 == 0:
                    board[r:r+cell, c:c+cell] = 255
        return board

    def _poses_to_enu(self, poses) -> np.ndarray:
        if not poses:
            return np.zeros((0, 3))
        origin = poses[0]
        arr = np.zeros((len(poses), 3))
        for i, p in enumerate(poses):
            dlat_m = (p.lat - origin.lat) * (np.pi / 180.0) * EARTH_R
            dlon_m = (
                (p.lon - origin.lon) * (np.pi / 180.0)
                * EARTH_R * np.cos(np.radians(origin.lat))
            )
            # ENU: x=east, y=north, z=up
            arr[i] = [dlon_m, dlat_m, p.alt - origin.alt]
        return arr

    def _estimates_to_enu(self, origin) -> np.ndarray:
        if origin is None:
            return np.zeros((0, 3))
        rows = []
        for _, est in self._estimates:
            if est is None:
                continue
            lat, lon, alt = est
            dlat_m = (lat - origin.lat) * (np.pi / 180.0) * EARTH_R
            dlon_m = (
                (lon - origin.lon) * (np.pi / 180.0)
                * EARTH_R * np.cos(np.radians(origin.lat))
            )
            rows.append([dlon_m, dlat_m, alt])
        return np.array(rows) if rows else np.zeros((0, 3))
