"""Input adapters for receiving distance measurements from various sources."""

import json
import logging
import socket
import struct
import threading
import time
from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, Optional

import serial

logger = logging.getLogger(__name__)


class InputAdapter(ABC):
    """Base class for input adapters."""

    @abstractmethod
    def start(self, callback: Callable[[Dict[int, Dict[str, float]]], None]):
        """Start receiving data and call callback with distance measurements.

        callback(distances: Dict[int, Dict[str, float]])
            Contains distance measurements for each anchor.
            distances[anchor_id] = {"distance": distance, "timestamp": timestamp}
        """
        pass

    @abstractmethod
    def stop(self):
        """Stop receiving data."""
        pass

    @abstractmethod
    def is_running(self) -> bool:
        """Check if adapter is running."""
        pass


class SerialInputAdapter(InputAdapter):
    """Serial port input adapter (from RTK-DW1000 implementation)."""

    def __init__(self, port: str, baud: int = 460800, timeout: float = 1.0):
        """Initialize serial input adapter."""
        self.port = port
        self.baud = baud
        self.timeout = timeout
        self.ser: serial.Serial | None = None
        self.running: bool = False
        self.callback: Callable[[Dict[int, Dict[str, float]]], None] | None = None
        self.thread: threading.Thread | None = None
        self.buffer: CircularBuffer = CircularBuffer(100, element_size=7)

    def start(self, callback: Callable[[Dict[int, Dict[str, float]]], None] | None = None):
        """Start reading from serial port."""
        if callback is not None:
            self.callback = callback
        try:
            self.ser = serial.Serial(
                self.port,
                self.baud,
                timeout=self.timeout,
                stopbits=serial.STOPBITS_ONE,
                bytesize=serial.EIGHTBITS,
                parity="E",
            )
            self.running = True
            self.thread = threading.Thread(target=self._read_loop, daemon=True)
            self.thread.start()
            logger.info("Serial adapter started on %s at %s baud", self.port, self.baud)
        except Exception as e:
            logger.error("Failed to connect to serial port: %s", e)
            raise

    def stop(self):
        """Stop reading from serial port."""
        self.running = False
        if self.ser:
            self.ser.close()
        if self.thread:
            self.thread.join(timeout=2.0)

    def is_running(self) -> bool:
        """Check if adapter is running."""
        return self.running

    def _read_loop(self):
        """Read from serial adapter."""
        ranges = {}
        while self.running:
            try:
                if self.ser.in_waiting > 0:
                    response = self.ser.read_all()
                    self.buffer.append(response)

                    while self.buffer.size > 0:
                        msg = self.buffer.pop()
                        if msg is None:
                            continue

                        message = Message(msg)
                        anchor_id = message.id
                        raw_val = message.data / 1000.0  # Convert mm to meters

                        if anchor_id is not None:
                            ranges[anchor_id] = {"distance": raw_val, "timestamp": time.time_ns() // 1000}

                    if ranges and self.callback:
                        self.callback(ranges.copy())
                        ranges.clear()

                time.sleep(0.01)  # Small delay to prevent CPU spinning
            except Exception as e:
                logger.error("Error in serial read loop: %s", e)
                time.sleep(0.1)


class MQTTInputAdapter(InputAdapter):
    """MQTT input adapter."""

    def __init__(self, broker: str, port: int = 1883, topic: str = "uwb/distances"):
        """Initialize MQTT input adapter."""
        import paho.mqtt.client as mqtt

        self.mqtt = mqtt
        self.broker = broker
        self.port = port
        self.topic = topic
        self.client: mqtt.Client | None = None
        self.running = False
        self.callback: Callable[[Dict[int, Dict[str, float]]], None] | None = None

    def start(self, callback: Callable[[Dict[int, Dict[str, float]]], None]):
        """Start MQTT subscriber."""
        try:
            self.callback = callback
            self.client = self.mqtt.Client()
            self.client.on_connect = self._on_connect
            self.client.on_message = self._on_message
            self.client.connect(self.broker, self.port, 60)
            self.running = True
            self.client.loop_start()
            logger.info(
                "MQTT adapter started on %s:%s, topic: %s",
                self.broker,
                self.port,
                self.topic,
            )
        except ImportError:
            raise ImportError("paho-mqtt is required for MQTT adapter. Install with: pip install paho-mqtt")
        except Exception as e:
            logger.error("Failed to connect to MQTT broker: %s", e)
            raise

    def stop(self):
        """Stop MQTT subscriber."""
        self.running = False
        if self.client:
            self.client.loop_stop()
            self.client.disconnect()

    def is_running(self) -> bool:
        """Check if adapter is running."""
        return self.running

    def _on_connect(self, client, userdata, flags, rc):
        """Handle MQTT connection callback."""
        del userdata, flags, rc
        client.subscribe(self.topic)

    def _on_message(self, client, userdata, msg):
        """MQTT message callback."""
        del client, userdata
        try:
            data = json.loads(msg.payload.decode())
            # Expected format: {"anchor_id": distance, ...} or [{"id": anchor_id, "distance": dist}, ...]
            if isinstance(data, dict):
                distances = {int(k): float(v) for k, v in data.items()}
            elif isinstance(data, list):
                distances = {int(item["id"]): float(item["distance"]) for item in data}
            else:
                logger.warning("Unexpected MQTT message format: %s", data)
                return

            if self.callback:
                self.callback(distances)
        except self.mqtt.exceptions.MqttError as e:
            logger.error("Error processing MQTT message: %s", e)


class UDPInputAdapter(InputAdapter):
    """UDP input adapter."""

    def __init__(self, host: str = "0.0.0.0", port: int = 5005):
        """Initialize UDP input adapter."""
        self.host = host
        self.port = port
        self.socket: socket.socket | None = None
        self.running = False
        self.callback: Callable[[Dict[int, Dict[str, float]]], None] | None = None
        self.thread: threading.Thread | None = None
        self.buffer = bytearray(1024)

    def start(self, callback: Callable[[Dict[int, Dict[str, float]]], None]):
        """Start UDP listener."""
        if callback is not None:
            self.callback = callback
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

            self.socket.bind((self.host, self.port))
            self.socket.setblocking(False)
            # self.ready = select.select([self.socket], [], [], 0.1)  # timeout
            self.socket.settimeout(1.0)
            self.running = True
            self.thread = threading.Thread(target=self._read_loop, daemon=True)
            self.thread.start()
            logger.info("UDP adapter started on %s:%s", self.host, self.port)
        except Exception as e:
            logger.error("Failed to start UDP adapter: %s", e)
            raise

    def stop(self):
        """Stop UDP listener."""
        self.running = False
        if self.socket:
            self.socket.close()
        if self.thread:
            self.thread.join(timeout=1.0)

    def is_running(self) -> bool:
        """Check if adapter is running."""
        return self.running

    def _read_loop(self):
        """Read from UDP adapter."""
        while self.running:
            if self.socket is None:
                self.running = False
                raise RuntimeError("Socket not initialized")

            data = self.socket.recv_into(self.buffer)
            if not data:
                self.running = False
                raise RuntimeError("Socket connection closed")
            if data == 0:
                raise RuntimeError("No data recieved")
            try:
                message = json.loads(self.buffer[:data].decode())
                # Expected format: {"anchor_id": distance, ...} or [{"id": anchor_id, "distance": dist}, ...]
                if isinstance(message, dict):
                    distances = {int(k): float(v) for k, v in message.items()}
                elif isinstance(message, list):
                    distances = {int(item["id"]): float(item["distance"]) for item in message}
                else:
                    continue

                if self.callback:
                    self.callback(distances)
                    return
            except json.JSONDecodeError as e:
                logger.warning("Invalid JSON in UDP message: %s", e)


class FileInputAdapter(InputAdapter):
    """File input adapter for reading from JSON file."""

    def __init__(self, filepath: str, poll_interval: float = 1.0):
        """Initialize file input adapter."""
        self.filepath = filepath
        self.poll_interval = poll_interval
        self.running = False
        self.callback: Callable[[Dict[int, Dict[str, float]]], None] | None = None
        self.thread: threading.Thread | None = None

    def start(self, callback: Callable[[Dict[int, Dict[str, float]]], None] | None = None):
        """Start reading from file."""
        if callback is not None:
            self.callback = callback
        self.running = True
        self.thread = threading.Thread(target=self._read_loop, daemon=True)
        self.thread.start()
        logger.info("File adapter started reading from %s", self.filepath)

    def stop(self):
        """Stop reading from file."""
        self.running = False
        if self.thread:
            self.thread.join(timeout=2.0)

    def is_running(self) -> bool:
        """Check if adapter is running."""
        return self.running

    def _read_loop(self):
        """Read from file adapter."""
        while self.running:
            try:
                with open(self.filepath, "r") as f:
                    data = json.load(f)
                    if isinstance(data, dict):
                        distances = {int(k): float(v) for k, v in data.items()}
                    elif isinstance(data, list):
                        distances = {int(item["id"]): float(item["distance"]) for item in data}
                    else:
                        continue

                    if self.callback:
                        self.callback(distances)
            except FileNotFoundError:
                logger.warning("File not found: %s", self.filepath)
            except json.JSONDecodeError as e:
                logger.warning("Invalid JSON in file: %s", e)
            except Exception as e:
                logger.error("Error reading file: %s", e)

            time.sleep(self.poll_interval)


# Helper classes from RTK-DW1000
class Message:
    """Message parser for serial protocol."""

    def __init__(self, data: bytes):
        """Initialize message parser."""
        self.id, self.data = struct.unpack("<BI", data)

    def __str__(self):
        """Return string representation of message."""
        return f"id: {self.id}, data: {self.data}"


class CircularBuffer:
    """Circular buffer for serial message parsing."""

    def __init__(self, size: int, element_size: int = 5):
        """Initialize circular buffer."""
        self.buffer = [b"\x00" for i in range(size)]
        self.head = 0
        self.size = 0
        self.tail = 0
        self.element_size = element_size

    def append(self, item: Optional[bytes]):
        """Append item to buffer."""
        if item is None:
            return
        items = item.split(b"\xff\xff\xff\x55")
        for item in items:
            if len(item) == 0:
                continue
            if len(item) > self.element_size:
                continue
            self.buffer[self.head] = item
            self.head = (self.head + 1) % len(self.buffer)
            self.size += 1
        if self.size > len(self.buffer):
            self.size = len(self.buffer)
            self.tail += 1

    def pop(self) -> Optional[bytes]:
        """Pop item from buffer."""
        if self.size == 0:
            return None
        item = self.buffer[self.tail]
        self.tail = (self.tail + 1) % len(self.buffer)
        self.size -= 1
        return item


INPUT_ADAPTERS = {
    "serial": SerialInputAdapter,
    "mqtt": MQTTInputAdapter,
    "udp": UDPInputAdapter,
    "file": FileInputAdapter,
}


def register_input_adapter(adapter_type: str, adapter: Any) -> None:
    """Register input adapter."""
    if adapter is None:
        raise ValueError("Adapter cannot be None")
    if not issubclass(adapter, InputAdapter):
        raise ValueError("Adapter must be a subclass of InputAdapter")
    if adapter_type in INPUT_ADAPTERS:
        raise ValueError(f"Adapter type {adapter_type} already registered")
    INPUT_ADAPTERS[adapter_type] = adapter
