"""Service for PX4 sensor calibration operations."""

from __future__ import annotations

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

from mavsdk import System
from mavsdk.calibration import CalibrationError

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

from ..models import CalibrationType
from ..px4_connection import get_system, mark_disconnected, _is_connection_loss
from ..utils.calibration_parser import (
    parse_accelerometer_progress,
    parse_gyroscope_progress,
    parse_magnetometer_progress,
    parse_horizon_progress,
)


def _format_grpc_error(exc: Exception) -> str:
    """Format gRPC errors into user-friendly messages."""
    error_msg = str(exc)
    
    if isinstance(exc, RpcError):
        # Check for COMMAND_DENIED (PX4 rejected the command)
        if "COMMAND_DENIED" in error_msg:
            return (
                "PX4 rejected the calibration command. "
                "This may occur if: (1) a calibration is already in progress, "
                "(2) the vehicle is armed or in flight, (3) the vehicle is not in the correct state. "
                "Please ensure the vehicle is disarmed, on the ground, and no other calibration is running."
            )
        # Check for common gRPC error patterns
        if "Connection reset by peer" in error_msg or "UNAVAILABLE" in error_msg:
            return (
                "Connection to PX4 device was lost. "
                "This may occur if the device disconnected, rebooted, or the connection timed out. "
                "Please check the device connection and try again."
            )
        elif "DEADLINE_EXCEEDED" in error_msg or "timeout" in error_msg.lower():
            return (
                "Connection timeout. The device may be unresponsive. "
                "Please check the device connection and try again."
            )
        elif "PERMISSION_DENIED" in error_msg:
            return "Permission denied. Please check device permissions."
        else:
            return f"gRPC error: {error_msg}"
    
    # For non-gRPC exceptions, return the error message as-is
    return error_msg


async def calibrate_sensors(
    port: str,
    baudrate: int,
    calibration_type: CalibrationType,
) -> AsyncGenerator[Dict[str, Any], None]:
    """
    Calibrate PX4 sensors and stream progress as structured dict payloads.

    This function is transport-agnostic: it does NOT produce SSE-formatted strings.
    It yields plain Python dicts suitable for JSON serialization. The API layer
    (FastAPI endpoint) is responsible for converting these dicts into SSE frames.

    The function properly handles connection cleanup and gRPC errors.
    """

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

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

        # Reuse a shared MAVSDK System() for this port/baudrate
        drone = await get_system(port=port, baudrate=baudrate)

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

        # Small delay to ensure connection is stable before starting calibration
        # This helps avoid COMMAND_DENIED errors from rapid-fire requests
        await asyncio.sleep(0.1)

        # Dispatch based on calibration type
        if calibration_type is CalibrationType.GYROSCOPE:
            # Gyroscope calibration
            yield {
                "type": "status",
                "message": "Starting gyroscope calibration...",
            }

            prev: Dict[str, Any] | None = None
            try:
                async for progress_data in drone.calibration.calibrate_gyro():
                    structured_data = parse_gyroscope_progress(progress_data, prev=prev)
                    prev = structured_data
                    yield structured_data
            except (RpcError, Exception) as exc:
                # Treat user-initiated cancellation as a non-error outcome
                if isinstance(exc, CalibrationError) and "CANCELLED" in str(exc):
                    yield {
                        "type": "calibration_cancelled",
                        "calibration_type": "gyroscope",
                        "message": "Gyroscope calibration cancelled.",
                    }
                    return

                yield {
                    "type": "error",
                    "message": f"Gyroscope calibration failed: {_format_grpc_error(exc)}",
                }
                raise

            yield {
                "type": "gyroscope_complete",
                "message": "Gyroscope calibration finished",
                "is_complete": True,
            }

        elif calibration_type is CalibrationType.ACCELEROMETER:
            # Accelerometer calibration
            yield {
                "type": "status",
                "message": "Starting accelerometer calibration...",
            }

            prev = None
            try:
                async for progress_data in drone.calibration.calibrate_accelerometer():
                    structured_data = parse_accelerometer_progress(progress_data, prev=prev)
                    prev = structured_data
                    yield structured_data
            except (RpcError, Exception) as exc:
                if isinstance(exc, CalibrationError) and "CANCELLED" in str(exc):
                    yield {
                        "type": "calibration_cancelled",
                        "calibration_type": "accelerometer",
                        "message": "Accelerometer calibration cancelled.",
                    }
                    return

                yield {
                    "type": "error",
                    "message": f"Accelerometer calibration failed: {_format_grpc_error(exc)}",
                }
                raise

            yield {
                "type": "accelerometer_complete",
                "message": "Accelerometer calibration finished",
                "is_complete": True,
            }

        elif calibration_type is CalibrationType.MAGNETOMETER:
            # Magnetometer calibration
            yield {
                "type": "status",
                "message": "Starting magnetometer calibration...",
            }

            prev = None
            try:
                async for progress_data in drone.calibration.calibrate_magnetometer():
                    structured_data = parse_magnetometer_progress(progress_data, prev=prev)
                    prev = structured_data
                    yield structured_data
            except (RpcError, Exception) as exc:
                if isinstance(exc, CalibrationError) and "CANCELLED" in str(exc):
                    yield {
                        "type": "calibration_cancelled",
                        "calibration_type": "magnetometer",
                        "message": "Magnetometer calibration cancelled.",
                    }
                    return

                yield {
                    "type": "error",
                    "message": f"Magnetometer calibration failed: {_format_grpc_error(exc)}",
                }
                raise

            yield {
                "type": "magnetometer_complete",
                "message": "Magnetometer calibration finished",
                "is_complete": True,
            }

        elif calibration_type is CalibrationType.HORIZON:
            # Level horizon calibration
            yield {
                "type": "status",
                "message": "Starting board level horizon calibration...",
            }

            prev = None
            try:
                async for progress_data in drone.calibration.calibrate_level_horizon():
                    structured_data = parse_horizon_progress(progress_data, prev=prev)
                    prev = structured_data
                    yield structured_data
            except (RpcError, Exception) as exc:
                if isinstance(exc, CalibrationError) and "CANCELLED" in str(exc):
                    yield {
                        "type": "calibration_cancelled",
                        "calibration_type": "horizon",
                        "message": "Level horizon calibration cancelled.",
                    }
                    return

                yield {
                    "type": "error",
                    "message": f"Horizon calibration failed: {_format_grpc_error(exc)}",
                }
                raise

            yield {
                "type": "horizon_complete",
                "message": "Level horizon calibration finished",
                "is_complete": True,
            }

        elif calibration_type is CalibrationType.ESC:
            # ESC calibration not implemented yet
            yield {
                "type": "status",
                "message": "Starting ESC calibration...",
            }
            yield {
                "type": "warning",
                "message": "ESC calibration not yet implemented",
            }
            yield {
                "type": "status",
                "message": "ESC calibration finished",
            }

    except (RpcError, Exception) as exc:
        # Handle connection and gRPC errors
        if _is_connection_loss(exc):
            await mark_disconnected(port, baudrate)
        error_message = _format_grpc_error(exc)
        yield {
            "type": "error",
            "message": f"Calibration failed: {error_message}",
        }
        raise

    finally:
        # With a shared System() we do not tear down the connection here.
        # Cleanup is handled at process level if needed.
        pass


async def cancel_calibration(
    port: str,
    baudrate: int,
) -> Dict[str, Any]:
    """
    Cancel any ongoing PX4 sensor calibration.

    This follows the recommended MAVSDK pattern:
    create a System, connect, then call `drone.calibration.cancel()`.
    """

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

    try:
        # Reuse the shared MAVSDK System() for this port/baudrate so that
        # cancel() is issued over the same underlying connection used by
        # calibration/start and other services.
        drone = await get_system(port=port, baudrate=baudrate)

        try:
            await drone.calibration.cancel()
        except CalibrationError as exc:
            # CalibrationError is specific to the calibration plugin
            return {
                "status": "error",
                "message": f"Calibration cancel failed: {exc}",
            }

        return {
            "status": "ok",
            "message": "Calibration cancelled.",
        }

    except (RpcError, Exception) as exc:
        if _is_connection_loss(exc):
            await mark_disconnected(port, baudrate)
        # Reuse the same error formatting logic as the main calibration flow
        return {
            "status": "error",
            "message": f"Calibration cancel failed: {_format_grpc_error(exc)}",
        }

    finally:
        # Shared System is kept alive; no per-call cleanup required here.
        pass
