"""
Windsuite SDK main module.
"""

from __future__ import annotations

import asyncio
import datetime
import math
import traceback
from collections.abc import Awaitable, Callable
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Protocol
from urllib.parse import urlparse

import requests
import socketio  # type: ignore[import-untyped]
from loguru import logger

from ._models_internal import TrackingDataDict, dict_to_dataclass
from .fan_control import FanControlBuilder, MachineLayout
from .models import FullTrackedWindProbeData, ModuleInfo, ModuleType, TrackedProbeData, TrackingData, Vec3

logger.remove()


class CallbackWindtrackData(Protocol):
    """
    This represents a callback function called whenever new windtrack data is received.

    It will pass the received data for the user to use.

    Args:
        data (dict[str, TrackingData]): Dictionary mapping object names (str) to their tracking data (TrackingData)

    """

    def __call__(self, data: dict[str, TrackingData]) -> None: ...


class CallbackWindprobeData(Protocol):
    """
    This represents a callback function called whenever new windprobe data is received.

    It will pass the received data for the user to use.

    Args:
        data (WindProbeData): The windprobe data received

    """

    def __call__(self, data: WindProbeData) -> None: ...


class CallbackTrackedProbeData(Protocol):
    """
    This represents a callback function called whenever new tracked windprobe data is received.

    It will pass the received data for the user to use.

    Args:
        data (TrackedProbeData): The tracked windprobe data received

    Note:
    - This callback is separate from CallbackWindprobeData to allow users to choose
    whether they want just the windprobe data or the full tracked data.

    - It also gives the user access to the best windprobe data if available depending on tracking (See TrackedProbeData comment)

    """

    def __call__(self, data: TrackedProbeData) -> None: ...


class CallbackModuleUpdate(Protocol):
    """
    This represents a callback function called whenever new module update data is received.

    It will pass the received data for the user to use.

    Args:
        data (dict[tuple[int, int], ModuleInfo]): Dictionary mapping module positions (row, col) to their ModuleInfo

    """

    def __call__(self, data: dict[tuple[int, int], ModuleInfo]) -> None: ...


class WindsuiteSDK:
    def __init__(
        self,
        base_url: str,
        port_rest_api: int = 8000,
        sio_port: int = 8001,
        sio_port_rt: int = 8002,
    ) -> None:
        self.port_rest_api = port_rest_api
        self.sio_port = sio_port
        self.sio_port_rt = sio_port_rt

        base_url = self.__parse_base_url(base_url)

        self.rest_base_url = base_url + f":{self.port_rest_api}" + "/api"
        self.sio_base_url = base_url + f":{self.sio_port}"
        self.sio_rt_base_url = base_url + f":{self.sio_port_rt}"
        self.__check_rest_api(base_url)

        self.__callback_windprobe: CallbackWindprobeData | None = None
        self.__callback_tracked_probe_data: CallbackTrackedProbeData | None = None
        self.__callback_windtrack: CallbackWindtrackData | None = None
        self.__callback_module_update: CallbackModuleUpdate | None = None

        self._executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="windsuite")
        self._stop_events: list[asyncio.Event] = []
        self._loops: list[asyncio.AbstractEventLoop] = []

        self.__sio: socketio.AsyncClient
        self.__sio_rt: socketio.AsyncClient

        self.current_layout = MachineLayout.from_api(self.rest_base_url)
        self.fan_controller = FanControlBuilder(self.current_layout)

    def __parse_base_url(self, url: str) -> str:
        """
        Parse the base URL and ensure it has the correct format
        If the given URL is missing the http:// prefix, it will be added automatically.
        If the URL has a trailing slash, it will be removed.
        If the URL contains a port, it will be removed
        An error will be raised if the URL is not valid.

        Args:
            url (str): The base URL to parse (e.g., "192.168.88.40" or "http://192.168.88.40")

        Returns:
            str: The parsed base URL in the format "http://<IP_ADDRESS>"
        Raises:
            ValueError: If the URL format is invalid

        """

        # Final check to ensure URL is valid
        def _raise_invalid_url(msg: str = ""):
            error_msg = f"Invalid URL format: {url}. {msg}Please provide a valid URL (e.g., 'http://192.168.88.40')."
            logger.error(error_msg)
            raise ValueError(error_msg)

        if not url.startswith("http://"):
            url = "http://" + url

        # Remove if there is multiple http:// in the URL (e.g., if user provided "http://)
        url = url.replace("http://http://", "http://")
        url = url.removesuffix("/")

        # Remove port and everything after it
        url = url.split(":")[0] + "://" + url.split("://")[1].split("/")[0]

        result = urlparse(url)
        if not all([result.scheme, result.netloc]):
            _raise_invalid_url()

        # Check for absurd IP like 172.16.20.21212 (last octet > 255 or not 4 octets)
        try:
            ip = result.netloc
            if ip.count(".") == 3:
                octets = ip.split(".")
            if len(octets) != 4 or not all(o.isdigit() for o in octets):
                _raise_invalid_url("Malformed IP address. ")
            if any(int(o) < 0 or int(o) > 255 for o in octets):
                _raise_invalid_url("IP octet out of range (0-255). ")
            # Optionally, could check for valid hostnames here too
        except ValueError:
            _raise_invalid_url("Malformed IP address. ")

        return url

    def __check_rest_api(self, base_url: str) -> bool:
        """
        Check if the REST API is reachable at the given base URL.
        Raises a RuntimeError with a descriptive message if the connection fails.
        """
        url = self.rest_base_url + "/configurator/health"
        headers = {"accept": "application/json"}
        try:
            response = requests.get(url, headers=headers, timeout=5)
            response.raise_for_status()
            return True
        except requests.RequestException as e:
            error_msg = (
                f"Could not connect to Windsuite server at {url}.\n"
                f"Reason: {e}\n"
                "Please check that the server is running and reachable at the specified address."
            )
            logger.error(error_msg)
            raise RuntimeError(error_msg) from e

    def __start_async_thread(
        self,
        async_routine: Callable[[asyncio.Event], Awaitable[None]],
        thread_name: str,
    ) -> None:
        """
        Helper to start an async routine in a background thread using ThreadPoolExecutor.

        Args:
            async_routine: The async coroutine function to run (receives a stop_event)
            thread_name: Name for the thread

        """

        def run_async() -> None:
            loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop=loop)
            stop_event = asyncio.Event()

            # Store the loop and event for cleanup
            self._loops.append(loop)
            self._stop_events.append(stop_event)

            try:
                loop.run_until_complete(async_routine(stop_event))
            finally:
                loop.close()

        self._executor.submit(run_async)

    def start_communication(self) -> None:
        """
        Called to start real-time communication with the Windsuite server

        ---

        When calling this:
        - Windprobe data will be received and passed to the registered windprobe callback
        - Windtrack data will be received and passed to the registered windtrack callback
        - Module updates will be received and passed to the registered module update callback
        - Facility updates (layout changes) will be handled automatically
        """
        self.__start_async_thread(self.__socketio_client_routine, "socketio_client_thread")
        self.__start_async_thread(self.__socketio_rt_client_routine, "socketio_rt_client_thread")

    def cleanup(self) -> None:
        """
        Stops all communication threads and disconnects from socketio servers.
        """
        logger.info("Starting cleanup...")

        for loop, stop_event in zip(self._loops, self._stop_events, strict=True):
            loop.call_soon_threadsafe(stop_event.set)

        logger.info("Shutting down thread pool...")
        self._executor.shutdown(wait=True, cancel_futures=False)

        logger.info("Cleanup complete.")

    def zero_windprobe(self) -> bool:
        """
        Trigger the zeroing of the windprobe
        """
        try:
            response = requests.post(f"{self.rest_base_url}/windprobe/zero", timeout=5)
            response.raise_for_status()
            return True
        except requests.RequestException as e:
            logger.error(f"Error zeroing windprobe: {e}")
            return False

    def set_psu(self, state: bool) -> bool:
        """
        Set the PSU of all modules to ON or OFF.
        """
        try:
            response = requests.patch(
                f"{self.rest_base_url}/windcontrol/modules/psu?psu_state_to_set={str(state).lower()}", timeout=5
            )
            response.raise_for_status()
            return True
        except requests.RequestException as e:
            logger.error(f"Error setting PSU state: {e}")
            return False

    def start_acquisition(self, name: str) -> bool:
        """
        Start data acquisition on the server (if not already started).
        """
        try:
            payload = {"name": name, "started_at": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
            response = requests.post(f"{self.rest_base_url}/acquisitions/start", json=payload, timeout=5)
            response.raise_for_status()
            return True
        except requests.RequestException as e:
            logger.error(f"Error starting acquisition: {e}")
            return False

    def stop_acquisition(self) -> bool:
        """
        Stop data acquisition on the server.
        """
        try:
            payload = {"stopped_at": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
            response = requests.post(f"{self.rest_base_url}/acquisitions/stop", json=payload, timeout=5)
            response.raise_for_status()
            return True
        except requests.RequestException as e:
            logger.error(f"Error stopping acquisition: {e}")
            return False

    def __replace_sentinel_with_nan(self, vec: Vec3) -> None:
        """Replace -999 sentinel values with nan in a Vec3."""
        if vec.x == -999:
            vec.x = math.nan
        if vec.y == -999:
            vec.y = math.nan
        if vec.z == -999:
            vec.z = math.nan

    async def __handle_probe_data(self, data: dict[str, Any]) -> None:
        """Handle incoming probe data."""
        try:
            full_tracked_probe_data = dict_to_dataclass(FullTrackedWindProbeData, data)

            # ! WINDPROBE DATA
            # Replace -999 sentinel values with nan in wind velocity fields
            windprobe = full_tracked_probe_data.windprobe_data
            self.__replace_sentinel_with_nan(windprobe.wind_velocity_mps_probe_ref)
            self.__replace_sentinel_with_nan(windprobe.wind_velocity_mps_windshaper_ref)
            self.__replace_sentinel_with_nan(windprobe.wind_velocity_mps_windshaper_ref_corrected)

            if self.__callback_windprobe:
                self.__callback_windprobe(windprobe)

            # ! TRACKED PROBE DATA
            # ! UGLY STUFF, NEEDS REFACTORING ASAP
            if full_tracked_probe_data.curated_tracked_probe_data is not None:
                tracked_probe_data = TrackedProbeData(
                    timestamp_s=full_tracked_probe_data.windprobe_data.timestamp_s,
                    wind_velocity_mps=full_tracked_probe_data.curated_tracked_probe_data.best_windprobe_data.best_wind_velocity_mps,
                    position_data_mm=full_tracked_probe_data.curated_tracked_probe_data.probe_position_mm_windshaper_ref,
                    temperature_celcius=full_tracked_probe_data.curated_tracked_probe_data.best_windprobe_data.temperature_celcius,
                    atmospheric_pressure_hpascal=full_tracked_probe_data.curated_tracked_probe_data.best_windprobe_data.atmospheric_pressure_hpascal,
                    static_pressure_pascal=full_tracked_probe_data.curated_tracked_probe_data.best_windprobe_data.static_pressure_pascal,
                )
                if self.__callback_tracked_probe_data:
                    self.__callback_tracked_probe_data(tracked_probe_data)
            # else:
            # print("No best probe data available for this tracked probe data.")

        except (KeyError, TypeError, ValueError) as e:
            logger.error(f"Error parsing probe data: {e}")
            logger.debug(f"Received data keys: {list(data.keys())}")

    async def __handle_tracking_data(self, data: dict[str, Any]) -> None:
        """Handle incoming tracking data."""
        try:
            windtrack_dict = dict_to_dataclass(TrackingDataDict, data)
            if self.__callback_windtrack:
                self.__callback_windtrack(windtrack_dict.data)
        except (KeyError, TypeError, ValueError) as e:
            logger.error(f"Error parsing tracking data: {e}")
            logger.debug(f"Received data keys: {list(data.keys())}")

    async def __handle_module_update(self, data: dict[str, Any]) -> None:
        """Handle incoming module update data."""
        try:
            modules_info: dict[tuple[int, int], ModuleInfo] = {}
            updated_modules = data.get("updated_modules", {})
            for mac, module_data in updated_modules.items():
                # Look up row/col from the layout using MAC address
                position = self.current_layout.get_module_by_mac(mac)
                if position is None:
                    logger.warning(f"Received update for unknown module MAC: {mac}")
                    continue
                row, col = position

                # ! 1 based indexing
                row += 1
                col += 1

                pwms_data: dict[str, list[float]] = module_data.get("pwms", {})
                rpms_data: dict[str, list[float]] = module_data.get("rpms", {})

                # Convert PWM values from 0-1000 range to 0-100 range
                target_pwm: list[list[float]] = [[value / 10.0 for value in pwms_data[key]] for key in sorted(pwms_data.keys())]
                if not target_pwm:
                    target_pwm = [[]]

                # Convert PWM values from 0-1000 range to 0-100 range
                current_pwm: list[list[float]] = [[value / 10.0 for value in pwms_data[key]] for key in sorted(pwms_data.keys())]
                if not current_pwm:
                    current_pwm = [[]]

                # Build current_rpm as list of lists (one inner list per layer)
                current_rpm: list[list[float]] = [rpms_data[key] for key in sorted(rpms_data.keys())]
                if not current_rpm:
                    current_rpm = [[]]

                module_type_str = module_data["type"]
                if module_type_str in {e.value for e in ModuleType}:
                    module_type_str = ModuleType(module_type_str).value
                else:
                    logger.warning("Unknown module type received: {}", module_type_str)

                module_info = ModuleInfo(
                    row=row,
                    col=col,
                    mac=module_data["mac"],
                    ip=module_data["ip"],
                    type=module_type_str,
                    lifepoints=module_data.get("lifepoints"),
                    target_pwm=target_pwm,
                    current_pwm=current_pwm,
                    current_rpm=current_rpm,
                    target_psu_state=bool(module_data.get("psu_command")),
                    current_psu_state=bool(module_data.get("is_psu")),
                    is_connected=bool(module_data.get("is_connected")),
                    # ! Optional ESC and PSU data
                    esc_rpm=module_data.get("esc_rpm"),
                    applied_pwm_us=module_data.get("applied_pwm_us"),
                    esc_voltage_mv=module_data.get("esc_voltage_mv"),
                    esc_bus_current_ma=module_data.get("esc_bus_current_ma"),
                    esc_motor_current_ma=module_data.get("esc_motor_current_ma"),
                    esc_index=module_data.get("esc_index"),
                    esc_temp_mos_celcius=module_data.get("esc_temp_mos_celcius"),
                    esc_temp_capacitor_celcius=module_data.get("esc_temp_capacitor_celcius"),
                    esc_temp_mcu_celcius=module_data.get("esc_temp_mcu_celcius"),
                    esc_temp_motor_celcius=module_data.get("esc_temp_motor_celcius"),
                    esc_running_error=module_data.get("esc_running_error"),
                    esc_uptime=module_data.get("esc_uptime"),
                    applied_psu=module_data.get("applied_psu"),
                    psu_v_in=module_data.get("psu_v_in"),
                    psu_i_in=module_data.get("psu_i_in"),
                    psu_p_in=module_data.get("psu_p_in"),
                    psu_v_out=module_data.get("psu_v_out"),
                    psu_i_out=module_data.get("psu_i_out"),
                    psu_p_out=module_data.get("psu_p_out"),
                    psu_temp1=module_data.get("psu_temp1"),
                    psu_temp2=module_data.get("psu_temp2"),
                    psu_temp3=module_data.get("psu_temp3"),
                )
                modules_info[(row, col)] = module_info

            if self.__callback_module_update:
                self.__callback_module_update(modules_info)
            logger.debug("Parsed {} module(s): {}", len(modules_info), list(modules_info.keys()))
        except (KeyError, TypeError, ValueError) as e:
            logger.error(f"Error parsing module update data: {e} - {traceback.format_exc()}")
            logger.debug("Received data keys: {}", list(data.keys()))

    async def __handle_facility_update(self, data: dict[str, Any]) -> None:
        """Handle incoming facility update data (layout changes)."""
        try:
            windcontrol_layout = data.get("current_windcontrol_layout")
            if windcontrol_layout is None:
                logger.warning("Received facility_update without current_windcontrol_layout")
                return

            # Preserve the API URL from the current layout
            api_url = self.current_layout.api_url

            # Update the layout
            self.current_layout = MachineLayout.from_dict(windcontrol_layout, api_url=api_url)
            self.fan_controller = FanControlBuilder(self.current_layout)

            logger.info(
                f"Layout updated: {self.current_layout.nb_rows}x{self.current_layout.nb_columns} with {len(self.current_layout.modules)} modules"
            )
        except (KeyError, TypeError, ValueError) as e:
            logger.error(f"Error parsing facility update data: {e} - {traceback.format_exc()}")
            logger.debug("Received data keys: {}", list(data.keys()))

    async def __socketio_client_routine(self, stop_event: asyncio.Event) -> None:
        """
        Base socketio client for general events (windprobe data)
        """
        self.__sio: socketio.AsyncClient = socketio.AsyncClient(
            reconnection=False,
            request_timeout=5,
            reconnection_delay=1,
            reconnection_delay_max=5,
        )
        # ? Ignore rationales: sio is badly typed

        @self.__sio.event  # pyright: ignore[reportUnknownMemberType]
        async def connect() -> None:  # pyright: ignore[reportUnusedFunction]
            logger.info("Connected to SocketIO server.")

        @self.__sio.event  # pyright: ignore[reportUnknownMemberType]
        async def disconnect() -> None:  # pyright: ignore[reportUnusedFunction]
            logger.info("Disconnected from SocketIO server.")

        @self.__sio.event  # pyright: ignore[reportUnknownMemberType]
        async def probe_data(data: dict[str, Any]) -> None:  # pyright: ignore[reportUnusedFunction]
            await self.__handle_probe_data(data)

        @self.__sio.event  # pyright: ignore[reportUnknownMemberType]
        async def module_update(data: dict[str, Any]) -> None:  # pyright: ignore[reportUnusedFunction]
            await self.__handle_module_update(data)

        @self.__sio.event  # pyright: ignore[reportUnknownMemberType]
        async def facility_update(data: dict[str, Any]) -> None:  # pyright: ignore[reportUnusedFunction]
            await self.__handle_facility_update(data)

        try:
            await self.__sio.connect(  # pyright: ignore[reportUnknownMemberType]
                self.sio_base_url,
                transports=["websocket"],
                wait=True,
                wait_timeout=1,
            )
        except socketio.exceptions.ConnectionError as err:  # type: ignore[UnknownMemberType], Rationale : type hints of socketio are not strict enough
            logger.error(
                f"Socket IO connection error: {err}\nHINT:\n\t- Verify that the server is running\n\t- The URL is correct\n\t- You're on the same network as the server"
            )
            return

        await stop_event.wait()

        logger.info("Shutting down SIO client...")
        await self.__sio.disconnect()  # pyright: ignore[reportUnknownMemberType]
        logger.info("SIO client disconnected.")

    async def __socketio_rt_client_routine(self, stop_event: asyncio.Event) -> None:
        """
        Real-time socketio client
        """
        self.__sio_rt: socketio.AsyncClient = socketio.AsyncClient(
            reconnection=False,
            request_timeout=5,
            reconnection_delay=1,
            reconnection_delay_max=5,
        )
        # ? Ignore rationales: sio is badly typed

        @self.__sio_rt.event  # pyright: ignore[reportUnknownMemberType]
        async def connect() -> None:  # pyright: ignore[reportUnusedFunction]
            logger.info("Connected to SocketIO RT server.")

        @self.__sio_rt.event  # pyright: ignore[reportUnknownMemberType]
        async def disconnect() -> None:  # pyright: ignore[reportUnusedFunction]
            logger.info("Disconnected from SocketIO RT server.")

        @self.__sio_rt.event  # pyright: ignore[reportUnknownMemberType]
        async def tracking_data(data: dict[str, Any]) -> None:  # pyright: ignore[reportUnusedFunction]
            await self.__handle_tracking_data(data)

        try:
            await self.__sio_rt.connect(  # pyright: ignore[reportUnknownMemberType]
                self.sio_rt_base_url,
                transports=["websocket"],
                wait=True,
                wait_timeout=1,
            )
        except socketio.exceptions.ConnectionError as err:  # type: ignore[UnknownMemberType], Rationale : type hints of socketio are not strict enough
            logger.error(
                f"Socket IO RT connection error: {err}\nHINT:\n\t- Verify that the server is running\n\t- The URL is correct\n\t- You're on the same network as the server"
            )
            return

        await stop_event.wait()

        logger.info("Shutting down SIO RT client...")
        await self.__sio_rt.disconnect()  # pyright: ignore[reportUnknownMemberType]
        logger.info("SIO RT client disconnected.")

    def register_tracking_callback(
        self,
        callback: CallbackWindtrackData,
    ) -> None:
        """
        Register a callback for windtrack data.

        Args:
            callback (CallbackWindtrackData): The function to call when windtrack data is received.
                                         Receives a dict mapping object names to TrackingData.

        """
        if not callable(callback):
            raise TypeError("Callback must be a callable function.")

        self.__callback_windtrack = callback

    def register_windprobe_callback(
        self,
        callback: CallbackWindprobeData,
    ) -> None:
        """
        Register a callback for windprobe data.

        Args:
            callback (CallbackWindprobeData): The function to call when windprobe data is received.

        """
        if not callable(callback):
            raise TypeError("Callback must be a callable function.")

        self.__callback_windprobe = callback

    def register_tracked_probe_callback(
        self,
        callback: CallbackTrackedProbeData,
    ) -> None:
        """
        Register a callback for tracked windprobe data.

        Args:
            callback (CallbackTrackedProbeData): The function to call when tracked windprobe data is received.

        """
        if not callable(callback):
            raise TypeError("Callback must be a callable function.")

        self.__callback_tracked_probe_data = callback

    def register_module_update_callback(
        self,
        callback: CallbackModuleUpdate,
    ) -> None:
        """
        Register a callback for module update data.

        Args:
            callback (CallbackModuleUpdate): The function to call when module update data is received.

        """
        if not callable(callback):
            raise TypeError("Callback must be a callable function.")

        self.__callback_module_update = callback
