"""MAVLink I/O Bridge (Phase 4).

MAV-01: Sends GPS_INPUT (#233) over UART at 5–10 Hz via pymavlink.
MAV-02: Maps ESKF state + covariance → all GPS_INPUT fields.
MAV-03: Receives ATTITUDE / RAW_IMU, converts to IMUMeasurement, feeds ESKF.
MAV-04: Detects 3 consecutive frames with no position → sends NAMED_VALUE_FLOAT
        re-localisation request to ground station.
MAV-05: Telemetry at 1 Hz (confidence + drift) via NAMED_VALUE_FLOAT.

On dev/CI (pymavlink absent) every send/receive call silently no-ops via
MockMAVConnection so the rest of the pipeline remains testable.
"""

from __future__ import annotations

import asyncio
import logging
import math
import time
from typing import Callable, Optional

import numpy as np

from gps_denied.schemas import GPSPoint
from gps_denied.schemas.eskf import ConfidenceTier, ESKFState, IMUMeasurement
from gps_denied.schemas.mavlink import (
    GPSInputMessage,
    RelocalizationRequest,
    TelemetryMessage,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# pymavlink conditional import
# ---------------------------------------------------------------------------
try:
    from pymavlink import mavutil as _mavutil  # type: ignore
    _PYMAVLINK_AVAILABLE = True
    logger.info("pymavlink available — real MAVLink connection enabled")
except ImportError:
    _mavutil = None  # type: ignore
    _PYMAVLINK_AVAILABLE = False
    logger.info("pymavlink not available — using MockMAVConnection (dev/CI mode)")

# GPS epoch offset from Unix epoch (seconds)
_GPS_EPOCH_OFFSET = 315_964_800


# ---------------------------------------------------------------------------
# GPS time helpers (MAV-02)
# ---------------------------------------------------------------------------

def _unix_to_gps_time(unix_s: float) -> tuple[int, int]:
    """Convert Unix timestamp to (GPS_week, GPS_ms_of_week)."""
    gps_s = unix_s - _GPS_EPOCH_OFFSET
    gps_s = max(0.0, gps_s)
    week = int(gps_s // (7 * 86400))
    ms_of_week = int((gps_s % (7 * 86400)) * 1000)
    return week, ms_of_week


def _confidence_to_fix_type(confidence: ConfidenceTier) -> int:
    """Map ESKF confidence tier to GPS_INPUT fix_type (MAV-02)."""
    return {
        ConfidenceTier.HIGH: 3,    # 3D fix
        ConfidenceTier.MEDIUM: 3,  # 3D fix (VO tracking valid per solution.md)
        ConfidenceTier.LOW: 0,
        ConfidenceTier.FAILED: 0,
    }.get(confidence, 0)


def _eskf_to_gps_input(
    state: ESKFState,
    origin: GPSPoint,
    altitude_m: float = 0.0,
) -> GPSInputMessage:
    """Build a GPSInputMessage from ESKF state (MAV-02).

    Args:
        state:       Current ESKF nominal state.
        origin:      WGS84 ENU reference origin set at mission start.
        altitude_m:  Barometric altitude in metres MSL (from FC telemetry).
    """
    # ENU → WGS84
    east, north = state.position[0], state.position[1]
    cos_lat = math.cos(math.radians(origin.lat))
    lat_wgs84 = origin.lat + north / 111_319.5
    lon_wgs84 = origin.lon + east / (cos_lat * 111_319.5)

    # Velocity: ENU → NED
    vn = state.velocity[1]   # North = ENU[1]
    ve = state.velocity[0]   # East  = ENU[0]
    vd = -state.velocity[2]  # Down  = -Up

    # Accuracy from covariance (position block = rows 0-2, cols 0-2)
    cov_pos = state.covariance[:3, :3]
    sigma_h = math.sqrt(max(0.0, cov_pos[0, 0] + cov_pos[1, 1]))
    sigma_v = math.sqrt(max(0.0, cov_pos[2, 2]))
    speed_sigma = math.sqrt(max(0.0, (state.covariance[3, 3] + state.covariance[4, 4]) / 2.0))

    # Synthesised hdop/vdop (hdop ≈ σ_h / 5 maps to typical DOP scale)
    hdop = max(0.1, sigma_h / 5.0)
    vdop = max(0.1, sigma_v / 5.0)

    fix_type = _confidence_to_fix_type(state.confidence)

    now = state.timestamp if state.timestamp > 0 else time.time()
    week, week_ms = _unix_to_gps_time(now)

    return GPSInputMessage(
        time_usec=int(now * 1_000_000),
        time_week=week,
        time_week_ms=week_ms,
        fix_type=fix_type,
        lat=int(lat_wgs84 * 1e7),
        lon=int(lon_wgs84 * 1e7),
        alt=altitude_m,
        hdop=round(hdop, 2),
        vdop=round(vdop, 2),
        vn=round(vn, 4),
        ve=round(ve, 4),
        vd=round(vd, 4),
        speed_accuracy=round(speed_sigma, 2),
        horiz_accuracy=round(sigma_h, 2),
        vert_accuracy=round(sigma_v, 2),
        satellites_visible=10,
    )


# ---------------------------------------------------------------------------
# Mock MAVLink connection (dev/CI)
# ---------------------------------------------------------------------------

class MockMAVConnection:
    """No-op MAVLink connection used when pymavlink is not installed."""

    def __init__(self):
        self._sent: list[dict] = []
        self._rx_messages: list = []

    def mav(self):
        return self

    def gps_input_send(self, *args, **kwargs) -> None:  # noqa: D102
        self._sent.append({"type": "GPS_INPUT", "args": args, "kwargs": kwargs})

    def named_value_float_send(self, *args, **kwargs) -> None:  # noqa: D102
        self._sent.append({"type": "NAMED_VALUE_FLOAT", "args": args, "kwargs": kwargs})

    def recv_match(self, type=None, blocking=False, timeout=0.1):  # noqa: D102
        return None

    def close(self) -> None:
        pass


# ---------------------------------------------------------------------------
# MAVLinkBridge
# ---------------------------------------------------------------------------

class MAVLinkBridge:
    """Full MAVLink I/O bridge.

    Usage::

        bridge = MAVLinkBridge(connection_string="serial:/dev/ttyTHS1:57600")
        await bridge.start(origin_gps, eskf_instance)
        # ... flight ...
        await bridge.stop()
    """

    def __init__(
        self,
        connection_string: str = "udp:127.0.0.1:14550",
        output_hz: float = 5.0,
        telemetry_hz: float = 1.0,
        max_consecutive_failures: int = 3,
    ):
        self.connection_string = connection_string
        self.output_hz = output_hz
        self.telemetry_hz = telemetry_hz
        self.max_consecutive_failures = max_consecutive_failures

        self._conn = None
        self._origin: Optional[GPSPoint] = None
        self._altitude_m: float = 0.0

        # State shared between loops
        self._last_state: Optional[ESKFState] = None
        self._last_gps: Optional[GPSPoint] = None
        self._consecutive_failures: int = 0
        self._frames_since_sat: int = 0
        self._drift_estimate_m: float = 0.0

        # Callbacks
        self._on_imu: Optional[Callable[[IMUMeasurement], None]] = None
        self._on_reloc_request: Optional[Callable[[RelocalizationRequest], None]] = None

        # asyncio tasks
        self._tasks: list[asyncio.Task] = []
        self._running = False

        # Diagnostics
        self._sent_count: int = 0
        self._recv_imu_count: int = 0

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def set_imu_callback(self, cb: Callable[[IMUMeasurement], None]) -> None:
        """Register callback invoked for each received IMU packet (MAV-03)."""
        self._on_imu = cb

    def set_reloc_callback(self, cb: Callable[[RelocalizationRequest], None]) -> None:
        """Register callback invoked when re-localisation is requested (MAV-04)."""
        self._on_reloc_request = cb

    def update_state(self, state: ESKFState, altitude_m: float = 0.0) -> None:
        """Push a fresh ESKF state snapshot (called by processor per frame)."""
        self._last_state = state
        self._altitude_m = altitude_m
        if state.confidence in (ConfidenceTier.HIGH, ConfidenceTier.MEDIUM):
            # Position available
            self._consecutive_failures = 0
        else:
            self._consecutive_failures += 1

    def notify_satellite_correction(self) -> None:
        """Reset frames_since_sat counter after a satellite match."""
        self._frames_since_sat = 0

    def update_drift_estimate(self, drift_m: float) -> None:
        """Update running drift estimate (metres) for telemetry."""
        self._drift_estimate_m = drift_m

    async def start(self, origin: GPSPoint) -> None:
        """Open the connection and launch background I/O coroutines."""
        self._origin = origin
        self._running = True
        self._conn = self._open_connection()
        self._tasks = [
            asyncio.create_task(self._gps_output_loop(), name="mav_gps_output"),
            asyncio.create_task(self._imu_receive_loop(), name="mav_imu_input"),
            asyncio.create_task(self._telemetry_loop(), name="mav_telemetry"),
        ]
        logger.info("MAVLinkBridge started (conn=%s, %g Hz)", self.connection_string, self.output_hz)

    async def stop(self) -> None:
        """Cancel background tasks and close connection."""
        self._running = False
        for t in self._tasks:
            t.cancel()
        await asyncio.gather(*self._tasks, return_exceptions=True)
        self._tasks.clear()
        if self._conn:
            self._conn.close()
            self._conn = None
        logger.info("MAVLinkBridge stopped. sent=%d imu_rx=%d",
                    self._sent_count, self._recv_imu_count)

    def build_gps_input(self) -> Optional[GPSInputMessage]:
        """Build GPSInputMessage from current ESKF state (public, for testing)."""
        if self._last_state is None or self._origin is None:
            return None
        return _eskf_to_gps_input(self._last_state, self._origin, self._altitude_m)

    # ------------------------------------------------------------------
    # MAV-01/02: GPS_INPUT output loop
    # ------------------------------------------------------------------

    async def _gps_output_loop(self) -> None:
        """Send GPS_INPUT at output_hz.  MAV-01 / MAV-02."""
        interval = 1.0 / self.output_hz
        while self._running:
            try:
                msg = self.build_gps_input()
                if msg is not None:
                    self._send_gps_input(msg)
                    self._sent_count += 1

                    # MAV-04: check consecutive failures
                    if self._consecutive_failures >= self.max_consecutive_failures:
                        self._send_reloc_request()
            except Exception as exc:
                logger.warning("GPS output loop error: %s", exc)
            await asyncio.sleep(interval)

    def _send_gps_input(self, msg: GPSInputMessage) -> None:
        if self._conn is None:
            return
        try:
            if _PYMAVLINK_AVAILABLE and not isinstance(self._conn, MockMAVConnection):
                self._conn.mav.gps_input_send(
                    msg.time_usec,
                    msg.gps_id,
                    msg.ignore_flags,
                    msg.time_week_ms,
                    msg.time_week,
                    msg.fix_type,
                    msg.lat,
                    msg.lon,
                    msg.alt,
                    msg.hdop,
                    msg.vdop,
                    msg.vn,
                    msg.ve,
                    msg.vd,
                    msg.speed_accuracy,
                    msg.horiz_accuracy,
                    msg.vert_accuracy,
                    msg.satellites_visible,
                )
            else:
                # MockMAVConnection records the call
                self._conn.gps_input_send(
                    time_usec=msg.time_usec,
                    fix_type=msg.fix_type,
                    lat=msg.lat,
                    lon=msg.lon,
                )
        except Exception as exc:
            logger.error("Failed to send GPS_INPUT: %s", exc)

    # ------------------------------------------------------------------
    # MAV-03: IMU receive loop
    # ------------------------------------------------------------------

    async def _imu_receive_loop(self) -> None:
        """Receive ATTITUDE/RAW_IMU and invoke ESKF callback.  MAV-03."""
        while self._running:
            try:
                raw = self._recv_imu()
                if raw is not None:
                    self._recv_imu_count += 1
                    if self._on_imu:
                        self._on_imu(raw)
            except Exception as exc:
                logger.warning("IMU receive loop error: %s", exc)
            await asyncio.sleep(0.01)  # poll at ~100 Hz; blocks throttled by recv_match timeout

    def _recv_imu(self) -> Optional[IMUMeasurement]:
        """Try to read one IMU packet from the MAVLink connection."""
        if self._conn is None:
            return None
        if isinstance(self._conn, MockMAVConnection):
            return None  # mock produces no IMU traffic

        try:
            msg = self._conn.recv_match(type=["RAW_IMU", "SCALED_IMU2"], blocking=False, timeout=0.01)
            if msg is None:
                return None
            t = time.time()
            # RAW_IMU fields (all in milli-g / milli-rad/s — convert to SI)
            ax = getattr(msg, "xacc", 0) * 9.80665e-3   # milli-g → m/s²
            ay = getattr(msg, "yacc", 0) * 9.80665e-3
            az = getattr(msg, "zacc", 0) * 9.80665e-3
            gx = getattr(msg, "xgyro", 0) * 1e-3        # milli-rad/s → rad/s
            gy = getattr(msg, "ygyro", 0) * 1e-3
            gz = getattr(msg, "zgyro", 0) * 1e-3
            return IMUMeasurement(
                accel=np.array([ax, ay, az]),
                gyro=np.array([gx, gy, gz]),
                timestamp=t,
            )
        except Exception as exc:
            logger.debug("IMU recv error: %s", exc)
            return None

    # ------------------------------------------------------------------
    # MAV-04: Re-localisation request
    # ------------------------------------------------------------------

    def _send_reloc_request(self) -> None:
        """Send NAMED_VALUE_FLOAT re-localisation beacon (MAV-04)."""
        req = self._build_reloc_request()
        if self._on_reloc_request:
            self._on_reloc_request(req)
        if self._conn is None:
            return
        try:
            t_boot_ms = int((time.time() % (2**32 / 1000)) * 1000)
            for name, value in [
                ("RELOC_LAT", float(req.last_lat or 0.0)),
                ("RELOC_LON", float(req.last_lon or 0.0)),
                ("RELOC_UNC", float(req.uncertainty_m)),
            ]:
                if _PYMAVLINK_AVAILABLE and not isinstance(self._conn, MockMAVConnection):
                    self._conn.mav.named_value_float_send(
                        t_boot_ms,
                        name.encode()[:10],
                        value,
                    )
                else:
                    self._conn.named_value_float_send(time=t_boot_ms, name=name, value=value)
            logger.warning("Re-localisation request sent (failures=%d)", self._consecutive_failures)
        except Exception as exc:
            logger.error("Failed to send reloc request: %s", exc)

    def _build_reloc_request(self) -> RelocalizationRequest:
        last_lat, last_lon = None, None
        if self._last_state is not None and self._origin is not None:
            east = self._last_state.position[0]
            north = self._last_state.position[1]
            cos_lat = math.cos(math.radians(self._origin.lat))
            last_lat = self._origin.lat + north / 111_319.5
            last_lon = self._origin.lon + east / (cos_lat * 111_319.5)
            cov = self._last_state.covariance[:2, :2]
            sigma_h = math.sqrt(max(0.0, (cov[0, 0] + cov[1, 1]) / 2.0))
        else:
            sigma_h = 500.0
        return RelocalizationRequest(
            last_lat=last_lat,
            last_lon=last_lon,
            uncertainty_m=max(sigma_h * 3.0, 50.0),
            consecutive_failures=self._consecutive_failures,
        )

    # ------------------------------------------------------------------
    # MAV-05: Telemetry loop
    # ------------------------------------------------------------------

    async def _telemetry_loop(self) -> None:
        """Send confidence + drift at 1 Hz.  MAV-05."""
        interval = 1.0 / self.telemetry_hz
        while self._running:
            try:
                self._send_telemetry()
                self._frames_since_sat += 1
            except Exception as exc:
                logger.warning("Telemetry loop error: %s", exc)
            await asyncio.sleep(interval)

    def _send_telemetry(self) -> None:
        if self._last_state is None or self._conn is None:
            return

        fix_type = _confidence_to_fix_type(self._last_state.confidence)
        confidence_score = {
            ConfidenceTier.HIGH: 1.0,
            ConfidenceTier.MEDIUM: 0.6,
            ConfidenceTier.LOW: 0.2,
            ConfidenceTier.FAILED: 0.0,
        }.get(self._last_state.confidence, 0.0)

        telemetry = TelemetryMessage(
            confidence_score=confidence_score,
            drift_estimate_m=self._drift_estimate_m,
            fix_type=fix_type,
            frames_since_sat=self._frames_since_sat,
        )

        t_boot_ms = int((time.time() % (2**32 / 1000)) * 1000)
        for name, value in [
            ("CONF_SCORE", telemetry.confidence_score),
            ("DRIFT_M", telemetry.drift_estimate_m),
        ]:
            try:
                if _PYMAVLINK_AVAILABLE and not isinstance(self._conn, MockMAVConnection):
                    self._conn.mav.named_value_float_send(
                        t_boot_ms,
                        name.encode()[:10],
                        float(value),
                    )
                else:
                    self._conn.named_value_float_send(time=t_boot_ms, name=name, value=float(value))
            except Exception as exc:
                logger.debug("Telemetry send error: %s", exc)

    # ------------------------------------------------------------------
    # Connection factory
    # ------------------------------------------------------------------

    def _open_connection(self):
        if _PYMAVLINK_AVAILABLE:
            try:
                conn = _mavutil.mavlink_connection(self.connection_string)
                logger.info("MAVLink connection opened: %s", self.connection_string)
                return conn
            except Exception as exc:
                logger.warning("Cannot open MAVLink connection (%s) — using mock", exc)
        return MockMAVConnection()
