"""Service for PX4 shell connection management."""

import asyncio
from typing import Optional

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


class ShellConnection:
    """Manage a persistent MAVSDK shell connection with queued output."""

    def __init__(self, port: str, baudrate: int) -> None:
        self.port = port
        self.baudrate = baudrate
        self.drone: Optional[System] = None
        self.output_queue: "asyncio.Queue[str]" = asyncio.Queue()
        self.connected = False

    async def connect(self) -> None:
        """Establish connection to PX4 and start shell monitoring."""

        try:
            self.drone = System()
            serial = f"serial://{self.port}:{self.baudrate}"
            await self.drone.connect(system_address=serial)

            # Wait for connection to be established
            async for state in self.drone.core.connection_state():
                if state.is_connected:
                    self.connected = True
                    break

            asyncio.create_task(self._observe_shell())
        except (RpcError, Exception) as exc:
            self.connected = False
            if self.drone is not None:
                self.drone = None
            raise RuntimeError(f"Failed to connect to PX4 shell: {exc}") from exc

    async def _observe_shell(self) -> None:
        """Monitor shell output and place data into the queue."""

        if not self.drone:
            return

        try:
            async for output in self.drone.shell.receive():
                if not self.connected:
                    break
                await self.output_queue.put(output)
        except (RpcError, Exception) as exc:  # pragma: no cover - streaming loop
            error_msg = str(exc)
            if isinstance(exc, RpcError):
                if "Connection reset by peer" in error_msg or "UNAVAILABLE" in error_msg:
                    error_msg = "Connection to PX4 device was lost"
            await self.output_queue.put(f"Error receiving shell output: {error_msg}")
            self.connected = False

    async def send_command(self, command: str) -> None:
        """Send a command to the PX4 shell."""

        if not self.connected or not self.drone:
            raise RuntimeError("Not connected to PX4")
        try:
            await self.drone.shell.send(command)
        except (RpcError, Exception) as exc:
            self.connected = False
            raise RuntimeError(f"Failed to send command: {exc}") from exc

    async def receive_output(self) -> Optional[str]:
        """Receive shell output (non-blocking)."""

        try:
            return await asyncio.wait_for(self.output_queue.get(), timeout=0.1)
        except asyncio.TimeoutError:
            return None

    async def disconnect(self) -> None:
        """Close the connection and cleanup resources."""

        self.connected = False
        # Cleanup: Ensure the System object is properly released
        if self.drone is not None:
            self.drone = None

