"""Low-level servo controller on top of STS protocol."""

from __future__ import annotations

import threading
import time
from typing import Any

from .constants import (
    DEFAULT_BAUDRATE,
    DEFAULT_RETRIES,
    DEFAULT_TIMEOUT,
    MemAddr,
    StatusBit,
    TorqueEnable,
)
from .exceptions import CommunicationError, ConnectionError
from .port import SerialClient
from .protocol import (
    build_read_frame,
    build_sync_position_frame,
    build_sync_read_frame,
    build_write_frame,
    bytes_to_int16,
    parse_response_frame,
    parse_sync_read_responses_with_errors,
)


class ServoController:
    """Thread-safe serial controller for STS register access."""

    def __init__(
        self,
        port: str,
        baudrate: int = DEFAULT_BAUDRATE,
        timeout: float = DEFAULT_TIMEOUT,
        retries: int = DEFAULT_RETRIES,
        debug: bool = False,
    ) -> None:
        self._client = SerialClient(port=port, baudrate=baudrate, timeout=timeout)
        self._retries = max(0, retries)
        self._debug = debug
        self._lock = threading.Lock()
        self._error_lock = threading.Lock()
        self._last_error_codes: dict[int, int] = {}
        self._stats_lock = threading.Lock()
        self._connected_since: float | None = None
        self._connected_seconds_total: float = 0.0
        self._no_response_count: int = 0

    @property
    def port(self) -> str:
        """Return the configured serial port path."""
        return self._client.port

    def connect(self) -> bool:
        """Open serial link."""
        self._client.open()
        with self._stats_lock:
            if self._connected_since is None:
                self._connected_since = time.perf_counter()
        return True

    def disconnect(self, disable_ids: list[int] | None = None) -> None:
        """Disable requested servos then close serial link."""
        if disable_ids:
            for servo_id in disable_ids:
                try:
                    self.write_register(
                        servo_id, MemAddr.TORQUE_ENABLE, [TorqueEnable.OFF]
                    )
                except Exception:
                    if self._debug:
                        print(f"[debug] torque disable failed for servo {servo_id}")
        # Acquire _lock to ensure no concurrent _txrx/read_registers_sync
        # is in progress before closing the serial port.
        with self._lock:
            with self._stats_lock:
                if self._connected_since is not None:
                    self._connected_seconds_total += (
                        time.perf_counter() - self._connected_since
                    )
                    self._connected_since = None
            self._client.close()

    def is_connected(self) -> bool:
        """Return connection state."""
        return self._client.is_open

    def _txrx(
        self, frame: bytes, expect_response: bool = True, read_size: int = 64
    ) -> tuple[int, int, list[int]] | None:
        if not self.is_connected():
            raise ConnectionError("Not connected")

        last_exc: Exception | None = None
        for _ in range(self._retries + 1):
            try:
                with self._lock:
                    if self._debug:
                        print(f"TX: {frame.hex()}")

                    self._client.reset_input_buffer()
                    self._client.write(frame)
                    self._client.flush()

                    if not expect_response:
                        return None

                    response = self._client.read_available(
                        read_size,
                        poll_timeout=self._client.calc_poll_timeout(read_size),
                    )

                    if self._debug:
                        print(f"RX: {response.hex()}")

                if not response:
                    raise CommunicationError("No response received")

                parsed = parse_response_frame(response)
                if parsed is None:
                    raise CommunicationError(f"Invalid response: {response.hex()}")
                servo_id, error_code, _ = parsed
                self._cache_error_code(servo_id, error_code)
                return parsed
            except Exception as exc:
                last_exc = exc
                if self._debug:
                    print(f"[debug] _txrx retry: {exc}")
                time.sleep(0.002)

        if (
            isinstance(last_exc, CommunicationError)
            and str(last_exc) == "No response received"
        ):
            self._record_no_response(1)
        raise CommunicationError(str(last_exc) if last_exc else "Communication failed")

    def _send_instruction(self, frame: bytes, read_size: int = 6) -> bool:
        """Send a pre-built instruction frame, return True if error_code == 0.

        Internal helper for non-WRITE_DATA instructions
        (e.g. POSITION_CALIBRATE 0x0B).
        """
        result = self._txrx(frame, expect_response=True, read_size=read_size)
        if result is None:
            return False
        _, error_code, _ = result
        # _txrx already caches error_code; just check the value.
        return error_code == 0

    def read_register(self, servo_id: int, addr: int, length: int) -> list[int]:
        """Read register bytes from one servo.

        Only raises on errors that invalidate the data (ENCODER_ERROR).
        Other errors (voltage, temperature, etc.) are status warnings —
        the returned data is still valid and will be used.
        """
        frame = build_read_frame(servo_id, addr, length)
        result = self._txrx(frame, expect_response=True, read_size=6 + length)
        if result is None:
            raise CommunicationError("Read failed")
        _, error_code, params = result
        if error_code & StatusBit.READ_INVALIDATING_ERRORS:
            raise CommunicationError(
                f"Servo {servo_id} read error {error_code:#04x}: "
                "encoder data unreliable"
            )
        return params

    def get_cached_error_code(self, servo_id: int) -> int | None:
        """Get last response-frame error code for one servo."""
        with self._error_lock:
            return self._last_error_codes.get(servo_id)

    def get_cached_error_codes(self) -> dict[int, int]:
        """Get a snapshot of cached response-frame error codes."""
        with self._error_lock:
            return dict(self._last_error_codes)

    def write_register(self, servo_id: int, addr: int, data: list[int]) -> bool:
        """Write register bytes to one servo."""
        frame = build_write_frame(servo_id, addr, data)
        result = self._txrx(frame, expect_response=True, read_size=6)
        if result is None:
            return False
        _, error_code, _ = result
        return error_code == 0

    def read_registers_sync(
        self,
        servo_ids: list[int],
        addr: int,
        data_len: int,
    ) -> tuple[dict[int, list[int]], dict[int, int | None]]:
        """Generic sync read for arbitrary register.

        Returns (parsed_data, error_codes) where:
        - parsed_data: {servo_id: [byte, ...]} for servos that responded
        - error_codes: {servo_id: error_code or None}

        Servos that fail sync read are retried with individual reads.
        """
        if not servo_ids:
            return {}, {}

        frame = build_sync_read_frame(addr, data_len, servo_ids)
        expected = (data_len + 6) * len(servo_ids)
        poll_timeout = self._client.calc_poll_timeout(expected)
        with self._lock:
            self._client.reset_input_buffer()
            self._client.write(frame)
            self._client.flush()
            time.sleep(0.002 * len(servo_ids))
            raw = self._client.read_available(expected, poll_timeout=poll_timeout)

        parsed, errors = parse_sync_read_responses_with_errors(raw, servo_ids, data_len)
        for servo_id, error_code in errors.items():
            if error_code is not None:
                self._cache_error_code(servo_id, error_code)
        missed_ids = [
            servo_id for servo_id in servo_ids if parsed.get(servo_id) is None
        ]
        if missed_ids:
            self._record_no_response(len(missed_ids))

        # Fallback single read for missing servos.
        for servo_id in servo_ids:
            if parsed.get(servo_id) is None:
                try:
                    data = self.read_register(servo_id, addr, data_len)
                    parsed[servo_id] = data
                    errors[servo_id] = self.get_cached_error_code(servo_id)
                except Exception as exc:
                    if self._debug:
                        print(f"[debug] fallback read servo {servo_id}: {exc}")

        return parsed, errors

    def read_positions(self, servo_ids: list[int]) -> dict[int, int]:
        """Read positions for multiple servos via sync-read with fallback."""
        parsed, errors = self.read_registers_sync(
            servo_ids, MemAddr.PRESENT_POSITION, 2
        )
        result: dict[int, int] = {}
        for servo_id in servo_ids:
            payload = parsed.get(servo_id)
            err = errors.get(servo_id)
            if payload and len(payload) >= 2:
                if err is not None and err & StatusBit.READ_INVALIDATING_ERRORS:
                    result[servo_id] = 0
                else:
                    result[servo_id] = bytes_to_int16(payload[0], payload[1])
            else:
                result[servo_id] = 0
        return result

    def write_positions_sync(
        self,
        servo_positions: list[tuple[int, int]],
        speed: int,
        acceleration: int = 0,
    ) -> bool:
        """Sync write positions for multiple motors."""
        if not servo_positions:
            return True
        frame = build_sync_position_frame(
            servo_positions, speed=speed, acceleration=acceleration
        )
        self._txrx(frame, expect_response=False)
        return True

    def _cache_error_code(self, servo_id: int, error_code: int) -> None:
        """Update last error-code cache with a response-frame value."""
        with self._error_lock:
            self._last_error_codes[servo_id] = error_code

    def _record_no_response(self, count: int) -> None:
        """Accumulate no-response events."""
        with self._stats_lock:
            self._no_response_count += max(0, int(count))

    def get_comm_stats(self) -> dict[str, Any]:
        """Get communication stats for runtime monitoring."""
        with self._stats_lock:
            connected_seconds = self._connected_seconds_total
            if self._connected_since is not None:
                connected_seconds += time.perf_counter() - self._connected_since
            no_response_count = self._no_response_count
        miss_per_minute = (
            (no_response_count * 60.0 / connected_seconds)
            if connected_seconds > 0
            else 0.0
        )
        return {
            "no_response_count": no_response_count,
            "connected_duration_sec": round(connected_seconds, 3),
            "no_response_per_minute": round(miss_per_minute, 3),
        }

    def reset_comm_stats(self) -> None:
        """Reset communication stats counters."""
        with self._stats_lock:
            self._no_response_count = 0
            self._connected_seconds_total = 0.0
            self._connected_since = time.perf_counter() if self.is_connected() else None
