from __future__ import annotations

"""
Shared MAVSDK System() connection manager for the PX4 configuration API.

This module ensures that we only create and maintain a small number of
MAVSDK System instances (typically one per (port, baudrate)), so that:

- We avoid repeated mavsdk_server startup/teardown and port conflicts.
- Calibration start/cancel, health checks, and status_text streaming
  all operate on the same underlying PX4 connection.

Connection state is monitored by a background watcher per connection; when
the link is lost (connection_state reports not connected or gRPC error),
we set is_connected=False so the next get_system() creates a fresh System().
"""

import asyncio
import logging
from dataclasses import dataclass
from typing import Dict, Optional, Tuple

from mavsdk import System

try:
    from grpc import RpcError
except ImportError:  # pragma: no cover
    RpcError = Exception  # type: ignore[misc,assignment]

log = logging.getLogger(__name__)


def _is_connection_loss(exc: BaseException) -> bool:
    """Return True if the exception indicates connection loss (so we should mark disconnected)."""
    msg = str(exc)
    return (
        "UNAVAILABLE" in msg
        or "Socket closed" in msg
        or "Connection reset by peer" in msg
        or "StatusCode.UNAVAILABLE" in msg
    )


@dataclass
class _Px4Connection:
    system: System
    lock: asyncio.Lock
    is_connected: bool = False
    watcher_task: Optional[asyncio.Task] = None


_connections: Dict[Tuple[str, int], _Px4Connection] = {}
_connections_lock = asyncio.Lock()


async def _watcher_loop(port: str, baudrate: int, conn: _Px4Connection) -> None:
    """
    Background task: watch connection_state(); when not connected or on error,
    set conn.is_connected = False and exit so the next get_system() can reconnect.
    """
    try:
        async for state in conn.system.core.connection_state():
            if not state.is_connected:
                async with conn.lock:
                    conn.is_connected = False
                    conn.watcher_task = None
                log.info("Connection lost (connection_state) for %s at %s baud", port, baudrate)
                return
    except asyncio.CancelledError:
        raise
    except Exception as exc:
        async with conn.lock:
            conn.is_connected = False
            conn.watcher_task = None
        log.info(
            "Connection watcher error for %s at %s baud: %s",
            port,
            baudrate,
            exc,
            exc_info=False,
        )


async def mark_disconnected(port: str, baudrate: int) -> None:
    """
    Mark the connection for (port, baudrate) as disconnected: set is_connected=False
    and cancel the watcher task. Next get_system() will create a new System() and reconnect.
    """
    key = (port, baudrate)
    async with _connections_lock:
        conn = _connections.get(key)
        if conn is None:
            return
        async with conn.lock:
            conn.is_connected = False
            if conn.watcher_task:
                conn.watcher_task.cancel()
                conn.watcher_task = None
        log.info("Marked disconnected: %s at %s baud", port, baudrate)


async def is_connected(port: str, baudrate: int) -> bool:
    """
    Return the current in-memory connection state for (port, baudrate).
    No I/O and does not trigger reconnect. Used by the heartbeat route.
    """
    key = (port, baudrate)
    async with _connections_lock:
        conn = _connections.get(key)
        if conn is None:
            return False
        return conn.is_connected


async def cancel_all_watchers(shutdown_timeout_s: float = 2.0) -> None:
    """
    Cancel all connection watcher tasks so the process can exit cleanly.
    Call from FastAPI lifespan on shutdown.
    """
    async with _connections_lock:
        tasks = []
        for conn in _connections.values():
            if conn.watcher_task:
                conn.watcher_task.cancel()
                tasks.append(conn.watcher_task)
                conn.watcher_task = None
    for t in tasks:
        try:
            await asyncio.wait_for(asyncio.shield(t), timeout=shutdown_timeout_s)
        except (asyncio.CancelledError, asyncio.TimeoutError):
            pass


async def get_system(
    port: str,
    baudrate: int,
    timeout_s: float = 10.0,
) -> System:
    """
    Return a shared MAVSDK System instance for the given serial port/baudrate.

    The first caller for a (port, baudrate) pair will create and connect the
    System; subsequent callers reuse the same connection. If the connection
    is marked disconnected (watcher or mark_disconnected), the next caller
    gets a new System() and reconnects. If the connection cannot be established
    within `timeout_s`, an exception is raised.
    """

    key = (port, baudrate)

    # Ensure we have a connection record for this key
    async with _connections_lock:
        conn = _connections.get(key)
        if conn is None:
            conn = _Px4Connection(
                system=System(),
                lock=asyncio.Lock(),
                is_connected=False,
                watcher_task=None,
            )
            _connections[key] = conn

    # Only one task at a time should (re)connect this System
    async with conn.lock:
        if conn.is_connected:
            return conn.system

        # Reconnect path: cancel watcher if present, create new System
        if conn.watcher_task:
            conn.watcher_task.cancel()
            try:
                await conn.watcher_task
            except asyncio.CancelledError:
                pass
            conn.watcher_task = None

        new_system = System()
        conn.system = new_system
        serial = f"serial://{port}:{baudrate}"
        await conn.system.connect(system_address=serial)

        # Wait for connection state to report is_connected with a timeout
        start = asyncio.get_event_loop().time()
        connected = False

        try:
            async for state in conn.system.core.connection_state():
                if state.is_connected:
                    connected = True
                    break
                elapsed = asyncio.get_event_loop().time() - start
                if elapsed >= timeout_s:
                    break
        except (RpcError, Exception):
            connected = False

        if not connected:
            # Leave is_connected as False so a later caller can retry.
            raise TimeoutError(
                f"Timed out connecting to PX4 on {port} at {baudrate} baud."
            )

        conn.is_connected = True
        conn.watcher_task = asyncio.create_task(_watcher_loop(port, baudrate, conn))
        log.info("Connected to PX4 on %s at %s baud", port, baudrate)
        return conn.system

