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

from __future__ import annotations

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

from mavsdk import System

try:
    from grpc import RpcError
except ImportError:
    # Fallback if grpc is not directly available (shouldn't happen with mavsdk)
    RpcError = Exception

from ..config import settings
from ..px4_connection import get_system, mark_disconnected, _is_connection_loss
from ..models import (
    Px4HealthResponse,
    CalibrationHealth,
    HealthAndArmingCheckReport,
    HealthAndArmingCheckMode,
    HealthAndArmingCheckProblem,
    HealthComponentReport,
    BatteryStatus,
)


async def _get_stable_health(drone: System, duration_s: float = 2.0) -> Optional[Any]:
    """
    Stream health messages for the specified duration and return the latest.
    This skips initial stale messages (all False) and gets accurate values.
    """
    latest_health = None
    start_time = asyncio.get_event_loop().time()
    
    try:
        async for health in drone.telemetry.health():
            latest_health = health
            elapsed = asyncio.get_event_loop().time() - start_time
            if elapsed >= duration_s:
                break
    except Exception:
        pass
    
    return latest_health


async def _get_stable_gps_info(drone: System, duration_s: float = 2.0) -> Optional[Any]:
    """
    Stream GPS info messages for the specified duration and return the latest.
    This helps avoid transient/no-fix states right after connection.
    """
    latest_info = None
    start_time = asyncio.get_event_loop().time()

    try:
        async for info in drone.telemetry.gps_info():
            latest_info = info
            elapsed = asyncio.get_event_loop().time() - start_time
            if elapsed >= duration_s:
                break
    except Exception:
        pass

    return latest_info


async def _get_stable_battery(drone: System, duration_s: float = 2.0) -> Optional[Any]:
    """
    Stream battery messages for the specified duration and return the latest.
    This provides a small window to smooth out initial transients.
    """
    latest_battery = None
    start_time = asyncio.get_event_loop().time()

    try:
        async for battery in drone.telemetry.battery():
            latest_battery = battery
            elapsed = asyncio.get_event_loop().time() - start_time
            if elapsed >= duration_s:
                break
    except Exception:
        pass

    return latest_battery


def _convert_arming_checks_report(report: Any) -> HealthAndArmingCheckReport:
    """Convert MAVSDK HealthAndArmingCheckReport to our Pydantic model."""
    current = getattr(report, "current_mode_intention", None)
    current_mode = None
    if current is not None:
        problems = [
            HealthAndArmingCheckProblem(
                message=getattr(p, "message", ""),
                description=getattr(p, "description", ""),
                health_component=getattr(p, "health_component", ""),
            )
            for p in (getattr(current, "problems", None) or [])
        ]
        current_mode = HealthAndArmingCheckMode(
            mode_name=getattr(current, "mode_name", ""),
            can_arm_or_run=getattr(current, "can_arm_or_run", False),
            problems=problems,
        )

    health_components = [
        HealthComponentReport(
            name=getattr(c, "name", ""),
            label=getattr(c, "label", ""),
            is_present=getattr(c, "is_present", False),
            has_error=getattr(c, "has_error", False),
            has_warning=getattr(c, "has_warning", False),
        )
        for c in (getattr(report, "health_components", None) or [])
    ]

    all_problems = [
        HealthAndArmingCheckProblem(
            message=getattr(p, "message", ""),
            description=getattr(p, "description", ""),
            health_component=getattr(p, "health_component", ""),
        )
        for p in (getattr(report, "all_problems", None) or [])
    ]

    return HealthAndArmingCheckReport(
        current_mode_intention=current_mode,
        health_components=health_components,
        all_problems=all_problems,
    )


async def _get_arming_checks_report(drone: System) -> Optional[HealthAndArmingCheckReport]:
    """
    Fetch health and arming checks report from MAVSDK Events plugin.
    Returns None if Events plugin is not available (e.g. older PX4).
    """
    try:
        report = await drone.events.get_health_and_arming_checks_report()
        if report is None:
            return None
        return _convert_arming_checks_report(report)
    except Exception:
        return None


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 health snapshot using telemetry health flags."""

    port = port or settings.default_port
    baudrate = baudrate or settings.default_baudrate
    try:
        # Reuse a shared MAVSDK System() for this port/baudrate
        drone = await get_system(port=port, baudrate=baudrate, timeout_s=timeout_s)

        # In parallel, stream health/GPS/battery for ~2s and take the latest samples.
        # This keeps the overall latency close to 2s instead of 3x 2s sequentially.
        health_task = _get_stable_health(drone, duration_s=2.0)
        gps_task = _get_stable_gps_info(drone, duration_s=2.0)
        battery_task = _get_stable_battery(drone, duration_s=2.0)

        health_result, gps_result, battery_result = await asyncio.gather(
            asyncio.wait_for(health_task, timeout_s),
            asyncio.wait_for(gps_task, timeout_s),
            asyncio.wait_for(battery_task, timeout_s),
            return_exceptions=True,
        )

        # Health is required to build the response; if it failed or timed out, treat as disconnected.
        if isinstance(health_result, Exception) or health_result is None:
            return Px4HealthResponse(
                status="disconnected",
                connection_port=port,
                connection_baudrate=baudrate,
                calibration=None,
                gps_fix_type=None,
                satellites_used=None,
                armable=False,
                arming_checks_report=None,
                raw_health=None,
            )

        health = health_result

        # Use telemetry health flags directly
        calibration = CalibrationHealth(
            gyroscope_ok=health.is_gyrometer_calibration_ok,
            accelerometer_ok=health.is_accelerometer_calibration_ok,
            magnetometer_ok=health.is_magnetometer_calibration_ok,
        )

        # Determine status
        status = "ok" if (health.is_armable and not calibration.requires_calibration) else "degraded"

        # GPS snapshot for fix type and satellites used (optional)
        gps_info = None if isinstance(gps_result, Exception) else gps_result

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

        if gps_info is not None:
            fix_type = getattr(gps_info, "fix_type", None)
            num_satellites = getattr(gps_info, "num_satellites", None)

            if fix_type is not None:
                raw_name = getattr(fix_type, "name", str(fix_type))
                key = raw_name.upper()
                # Map MAVSDK enum names to the strings documented in Px4HealthResponse
                gps_fix_type_map = {
                    "NO_FIX": "no_fix",
                    "FIX_2D": "2d_fix",
                    "FIX_3D": "3d_fix",
                    "RTK_FLOAT": "rtk_float",
                    "RTK_FIXED": "rtk_fix",
                }
                gps_fix_type = gps_fix_type_map.get(key, raw_name.lower())

            if isinstance(num_satellites, int):
                satellites_used = num_satellites

        # Battery snapshot for voltage / remaining percent / current (optional)
        battery_info = None if isinstance(battery_result, Exception) else battery_result

        battery_status: Optional[BatteryStatus] = None

        if battery_info is not None:
            voltage_v = getattr(battery_info, "voltage_v", None)
            remaining_raw = getattr(battery_info, "remaining_percent", None)
            current_a = getattr(battery_info, "current_battery_a", None)

            remaining_percent: Optional[float] = None
            if isinstance(remaining_raw, (int, float)):
                value = float(remaining_raw)
                # MAVSDK typically reports 0.0–1.0; map to 0–100 for convenience.
                remaining_percent = value * 100.0 if value <= 1.0 else value

            battery_status = BatteryStatus(
                voltage_v=float(voltage_v) if isinstance(voltage_v, (int, float)) else None,
                remaining_percent=remaining_percent,
                current_a=float(current_a) if isinstance(current_a, (int, float)) else None,
            )

        # Minimal raw health + GPS + battery for debugging
        raw_health: Dict[str, Any] = {
            "raw_health_message": str(health),
            "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_armable": health.is_armable,
            "gps_info": str(gps_info) if gps_info is not None else None,
            "battery_info": str(battery_info) if battery_info is not None else None,
        }

        # Fetch detailed health and arming checks from Events plugin (optional)
        arming_checks_report = await _get_arming_checks_report(drone)

        return Px4HealthResponse(
            status=status,
            connection_port=port,
            connection_baudrate=baudrate,
            calibration=calibration,
            gps_fix_type=gps_fix_type,
            satellites_used=satellites_used,
            armable=health.is_armable,
            arming_checks_report=arming_checks_report,
            battery=battery_status,
            raw_health=raw_health,
        )

    except (RpcError, Exception, TimeoutError) as exc:
        # Return disconnected status on any connection/telemetry error
        if _is_connection_loss(exc):
            await mark_disconnected(port, baudrate)
        return Px4HealthResponse(
            status="disconnected",
            connection_port=port,
            connection_baudrate=baudrate,
            calibration=None,
            gps_fix_type=None,
            satellites_used=None,
            armable=False,
            arming_checks_report=None,
            raw_health=None,
        )


async def stream_status_text(
    port: str,
    baudrate: int,
) -> AsyncGenerator[Dict[str, Any], None]:
    """
    Stream PX4 status_text messages over MAVSDK.

    This is useful on setups where the Events plugin is unavailable and
    `health_and_arming_checks_report` cannot be used to explain arming
    failures. The consumer can inspect the text messages to understand
    why the vehicle cannot arm (e.g. battery over-voltage, GPS issues).
    """

    serial = f"serial://{port}:{baudrate}"

    try:
        # Initial status messages
        yield {
            "type": "status",
            "message": f"Connecting to {port}...",
        }

        # Reuse the shared MAVSDK System() so that status_text comes from the
        # same underlying connection as other services.
        drone = await get_system(port=port, baudrate=baudrate)

        yield {
            "type": "status",
            "message": "Connected to device",
        }

        # Stream status_text messages indefinitely until the caller disconnects
        async for msg in drone.telemetry.status_text():
            severity = None
            msg_type = getattr(msg, "type", None)
            if msg_type is not None:
                # MAVSDK enums usually expose a .name attribute
                severity = getattr(msg_type, "name", str(msg_type))

            text = getattr(msg, "text", "")

            yield {
                "type": "status_text",
                "severity": severity,
                "text": text,
            }

    except (RpcError, Exception) as exc:
        # Handle connection loss ("Socket closed") as a normal end of stream to
        # avoid noisy errors when PX4 or the link resets. Other errors are still
        # surfaced as an error message.
        message = str(exc)
        if "Socket closed" in message or "StatusCode.UNAVAILABLE" in message:
            await mark_disconnected(port, baudrate)
            yield {
                "type": "status",
                "message": "Disconnected from device (socket closed).",
            }
            # Do not re-raise: treat as graceful disconnect
            return

        yield {
            "type": "error",
            "message": f"Status text stream error: {exc}",
        }
        raise

    finally:
        # Shared System is managed by the connection manager; no teardown here.
        pass
