"""Service for PX4 health and calibration snapshot."""

from __future__ import annotations

import asyncio
from typing import Optional, Dict, Any, Iterable

from mavsdk import System

from ..config import settings
from ..models import Px4HealthResponse, CalibrationHealth


async def _get_health_once(drone: System, timeout_s: float = 5.0):
    """
    Wait for the first telemetry.health() message.

    Raises asyncio.TimeoutError if nothing is received within timeout_s.
    """
    async def _first():
        async for health in drone.telemetry.health():
            return health

    return await asyncio.wait_for(_first(), timeout_s)


async def _get_gps_info_once(drone: System, timeout_s: float = 5.0):
    """
    Wait for the first telemetry.gps_info() message.

    Raises asyncio.TimeoutError if nothing is received within timeout_s.
    """
    async def _first():
        async for gps_info in drone.telemetry.gps_info():
            return gps_info

    return await asyncio.wait_for(_first(), timeout_s)


SENSOR_INDEXES: Iterable[int] = (0, 1, 2)


async def _is_family_calibrated(
    drone: System,
    family: str,
    indexes: Iterable[int] = SENSOR_INDEXES,
) -> bool:
    """
    Param-based heuristic:

    A sensor family (gyro/accel/mag) is considered calibrated if ANY
    CAL_FAMILYn_ID != 0 for the probed indexes and the param read succeeds.
    """
    prefix_map = {
        "gyro": "CAL_GYRO",
        "accel": "CAL_ACC",
        "mag": "CAL_MAG",
    }
    prefix = prefix_map.get(family)
    if not prefix:
        return False

    for i in indexes:
        name = f"{prefix}{i}_ID"
        try:
            sensor_id = await drone.param.get_param_int(name)
        except Exception:
            # Ignore this index if param read fails
            continue
        if sensor_id != 0:
            return True

    return False


async def read_px4_health(
    port: Optional[str] = None,
    baudrate: Optional[int] = None,
    timeout_s: float = 10.0,
) -> Px4HealthResponse:
    """
    Connect to PX4 via MAVSDK and return a one-shot health snapshot.

    - Uses PX4 CAL_* params to derive calibration health.
    - Uses telemetry.health() for armable + positional flags.
    - Uses telemetry.gps_info() for fix type + satellites.
    - Returns Px4HealthResponse; never raises to the API layer.
    """

    # Fallback to global settings if not provided explicitly
    port = port or settings.default_port
    baudrate = baudrate or settings.default_baudrate
    serial = f"serial://{port}:{baudrate}"

    drone = System()

    try:
        # Connect to PX4
        await drone.connect(system_address=serial)

        # Optional: verify we actually got a connection state
        async for state in drone.core.connection_state():
            if not state.is_connected:
                # Treat as disconnected; don't bother with the rest
                return Px4HealthResponse(
                    status="disconnected",
                    connection_port=port,
                    connection_baudrate=baudrate,
                    calibration=None,
                    gps_fix_type=None,
                    satellites_used=None,
                    armable=False,
                    raw_health=None,
                )
            break

        # -------- PARAM-BASED CALIBRATION (source of truth) --------
        gyroscope_ok = await _is_family_calibrated(drone, "gyro")
        accelerometer_ok = await _is_family_calibrated(drone, "accel")
        magnetometer_ok = await _is_family_calibrated(drone, "mag")

        calibration = CalibrationHealth(
            gyroscope_ok=gyroscope_ok,
            accelerometer_ok=accelerometer_ok,
            magnetometer_ok=magnetometer_ok,
        )

        # -------- TELEMETRY HEALTH (for armable + debug) --------
        try:
            health = await _get_health_once(drone, timeout_s=timeout_s / 2)
        except asyncio.TimeoutError:
            health = None

        # -------- GPS INFO (unchanged) --------
        try:
            gps_info = await _get_gps_info_once(drone, timeout_s=timeout_s / 2)
        except asyncio.TimeoutError:
            gps_info = None

        gps_fix_type: Optional[str] = None
        satellites_used: Optional[int] = None

        if gps_info is not None:
            if hasattr(gps_info.fix_type, "name"):
                gps_fix_type = gps_info.fix_type.name.lower()
            else:
                gps_fix_type = str(gps_info.fix_type).lower()

            satellites_used = int(getattr(gps_info, "num_satellites", 0))

        # -------- STATUS + ARMABLE (respect your Literal) --------
        armable: Optional[bool] = None
        if health is not None:
            armable = health.is_armable

        # Map everything into your fixed status vocabulary:
        # - disconnected: handled in exception / early return
        # - ok: armable and all calibrated
        # - degraded: anything else (not armable or missing calibration)
        if armable and not calibration.requires_calibration:
            status = "ok"
        else:
            status = "degraded"

        # -------- RAW HEALTH FOR DEBUG/UI --------
        raw_health: Dict[str, Any] = {}

        # Telemetry-derived bits (if available)
        if health is not None:
            raw_health.update(
                {
                    "is_gyrometer_calibration_ok": health.is_gyrometer_calibration_ok,
                    "is_accelerometer_calibration_ok": health.is_accelerometer_calibration_ok,
                    "is_magnetometer_calibration_ok": health.is_magnetometer_calibration_ok,
                    "is_local_position_ok": getattr(health, "is_local_position_ok", None),
                    "is_global_position_ok": getattr(health, "is_global_position_ok", None),
                    "is_home_position_ok": getattr(health, "is_home_position_ok", None),
                    "is_armable": health.is_armable,
                }
            )
        else:
            raw_health.update(
                {
                    "is_gyrometer_calibration_ok": None,
                    "is_accelerometer_calibration_ok": None,
                    "is_magnetometer_calibration_ok": None,
                    "is_local_position_ok": None,
                    "is_global_position_ok": None,
                    "is_home_position_ok": None,
                    "is_armable": None,
                }
            )

        # Param-derived calibration snapshot as debug info (so you can see IDs in Swagger/UI later if you want)
        raw_health["calibration_params"] = {
            "source": "CAL_*_ID heuristic",
            "gyroscope_ok": gyroscope_ok,
            "accelerometer_ok": accelerometer_ok,
            "magnetometer_ok": magnetometer_ok,
        }

        return Px4HealthResponse(
            status=status,
            connection_port=port,
            connection_baudrate=baudrate,
            calibration=calibration,
            gps_fix_type=gps_fix_type,
            satellites_used=satellites_used,
            armable=armable,
            raw_health=raw_health,
        )

    except (asyncio.TimeoutError, Exception):
        # On any connection / telemetry failure, report "disconnected"
        return Px4HealthResponse(
            status="disconnected",
            connection_port=port,
            connection_baudrate=baudrate,
            calibration=None,
            gps_fix_type=None,
            satellites_used=None,
            armable=False,
            raw_health=None,
        )
