"""
Async JSON-RPC client with automatic request ID correlation for serial communication.

This module provides an async/await RPC client that automatically manages request IDs:
- Serial connection management via a SerialInterface abstraction
- JSON-RPC 2.0 command serialization with automatic request ID generation
- Internal response matching by ID (transparent to users)
- ID wraparound handling for uint32 range (0 to 4,294,967,295)
- Async/await API (mandatory)
- Timeout and retry logic
- Proper KeyboardInterrupt handling
- Boot output draining

Usage (IDs are managed internally, async/await required):
    async with RpcClient("/dev/ttyUSB0") as client:
        response = await client.send("ping")
        print(response.data)  # {"timestamp": ..., "uptimeMs": ...}
        print(response.success)  # True/False
        # Note: response._id exists but is internal - don't use it

    # With explicit serial interface (e.g., pyserial backend):
    from ci.util.serial_interface import create_serial_interface
    iface = create_serial_interface("/dev/ttyUSB0", use_pyserial=True)
    async with RpcClient("/dev/ttyUSB0", serial_interface=iface) as client:
        response = await client.send("ping")

    # Or with explicit connection management:
    client = RpcClient("/dev/ttyUSB0")
    await client.connect()
    await client.drain_boot_output()
    response = await client.send("testGpioConnection", [16, 17])
    await client.close()
"""

from __future__ import annotations

import asyncio
import json
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from ci.util.global_interrupt_handler import is_interrupted, notify_main_thread


if TYPE_CHECKING:
    from ci.util.crash_trace_decoder import CrashTraceDecoder
    from ci.util.serial_interface import SerialInterface


@dataclass
class RpcResponse:
    """Structured RPC response with internal request ID correlation.

    The request ID is managed internally for correlation and is not exposed
    in the public API. Users should only interact with success, data, and raw_line.
    """

    success: bool
    data: dict[str, Any]
    raw_line: str
    _id: int = 0  # Internal: Request ID from JSON-RPC 2.0 (always present, hidden from public API)

    def get(self, key: str, default: Any = None) -> Any:
        """Get value from response data."""
        return self.data.get(key, default)

    def __getitem__(self, key: str) -> Any:
        """Allow dict-like access to response data."""
        return self.data[key]

    def __contains__(self, key: str) -> bool:
        """Check if key exists in response data."""
        return key in self.data


class RpcTimeoutError(TimeoutError):
    """RPC operation timed out."""

    pass


class RpcError(Exception):
    """RPC operation failed."""

    pass


class RpcCrashError(RpcError):
    """Device crashed during RPC operation.

    Attributes:
        decoded_lines: Decoded stack trace lines (if decoding succeeded).
    """

    def __init__(self, message: str, decoded_lines: list[str] | None = None) -> None:
        super().__init__(message)
        self.decoded_lines = decoded_lines or []


class RpcClient:
    """Stateful JSON-RPC client for serial communication.

    Accepts a SerialInterface for serial I/O. Use create_serial_interface()
    to create the appropriate backend (pyserial or fbuild).

    If no serial_interface is provided, an fbuild backend is created by default.
    To use pyserial, create the interface externally and pass it in:

        from ci.util.serial_interface import create_serial_interface
        iface = create_serial_interface(port, use_pyserial=True)
        client = RpcClient(port, serial_interface=iface)

    Attributes:
        port: Serial port path
        baudrate: Serial baud rate (default: 115200)
        timeout: Default RPC response timeout in seconds
        response_prefix: Prefix for RPC responses (default: "REMOTE: ")
    """

    RESPONSE_PREFIX = "REMOTE: "

    # Post-ACK timeout floor for async (Phase 3) RPCs. After the device sends
    # `{"acknowledged": true}` the client restarts its wait with at least this
    # many seconds for the final response (or the caller-supplied timeout,
    # whichever is larger). Picked generously so a 100-LED runSingleTest on a
    # slow MCU has headroom; the typical run on Teensy 4 finishes in 80-90s.
    # See FastLED#3060.
    ASYNC_POST_ACK_TIMEOUT: float = 600.0

    def __init__(
        self,
        port: str,
        baudrate: int = 115200,
        timeout: float = 10.0,
        read_timeout: float = 0.5,
        serial_interface: SerialInterface | None = None,
        verbose: bool = False,
        crash_decoder: CrashTraceDecoder | None = None,
    ):
        """Initialize RPC client.

        Args:
            port: Serial port path (e.g., "/dev/ttyUSB0", "COM3")
            baudrate: Serial baud rate (default: 115200)
            timeout: Default timeout for RPC operations in seconds
            read_timeout: Timeout for individual serial reads (unused with fbuild)
            serial_interface: Pre-created SerialInterface. If None, an fbuild
                backend is created automatically on connect().
            verbose: If True, enable detailed RPC logging
            crash_decoder: Optional CrashTraceDecoder for inline crash decoding.
        """
        self.port = port
        self.baudrate = baudrate
        self.timeout = timeout
        self.read_timeout = read_timeout
        self.verbose = verbose
        self._serial: SerialInterface | None = serial_interface
        self._owns_serial = serial_interface is None
        self._connected = False  # Track whether connect() has been called
        self._next_id: int = (
            1  # Request ID counter for JSON-RPC 2.0 correlation (uint32 range)
        )
        self._crash_decoder: CrashTraceDecoder | None = crash_decoder

    @property
    def is_connected(self) -> bool:
        """Check if serial connection is open."""
        return self._serial is not None and self._connected

    async def connect(
        self,
        boot_wait: float = 3.0,
        drain_boot: bool = True,
        boot_signals: "list[str] | tuple[str, ...] | None" = None,
    ) -> None:
        """Open serial connection (async).

        Args:
            boot_wait: Maximum time to wait for device boot (seconds). When
                ``boot_signals`` is non-empty, this is the *deadline* for the
                spin-poll — the helper returns as soon as any signal appears,
                otherwise it falls back to waiting the full duration (matching
                the legacy fixed-sleep behavior so slow devices still work).
            drain_boot: Whether to drain boot output after connecting. Lines
                consumed while spinning for ``boot_signals`` count toward the
                drain, so this second pass only catches leftovers.
            boot_signals: Substrings to look for in serial output during the
                boot wait. When a match is found, ``connect()`` returns
                immediately instead of sleeping the full ``boot_wait``. Pass
                an empty list to opt out of spin-polling (legacy sleep path).
                Defaults to ``BOOT_SIGNALS_ANY`` which covers both ROM-loader
                and Arduino ``setup()`` banners.
        """
        if self._connected:
            return

        # Create serial interface if not provided externally (defaults to fbuild)
        if self._serial is None:
            from ci.util.serial_interface import create_serial_interface

            self._serial = create_serial_interface(
                port=self.port,
                baud_rate=self.baudrate,
            )
            self._owns_serial = True

        backend_name = type(self._serial).__name__
        if self.verbose:
            print(f"🔌 [RPC] Connecting to {self.port} @ {self.baudrate} baud...")
            print(f"🔌 [RPC] Using {backend_name} backend")

        await self._serial.connect()
        self._connected = True

        if self.verbose:
            print(f"✅ [RPC] Serial connection established")

        # Resolve default boot signals lazily so tests that don't touch
        # serial I/O don't import the util module.
        effective_signals: tuple[str, ...] | list[str]
        if boot_signals is None:
            from ci.util.boot_wait import BOOT_SIGNALS_ANY

            effective_signals = BOOT_SIGNALS_ANY
        else:
            effective_signals = boot_signals

        poll_drained_lines: list[str] = []

        if boot_wait > 0:
            if self.verbose:
                print(f"⏳ [RPC] Waiting up to {boot_wait}s for device boot...")
            from ci.util.boot_wait import wait_for_signal

            def _mirror(line: str) -> None:
                if self.verbose:
                    print(f"  [boot-poll] {line[:120]}")

            # Spin-poll for a boot signal. If none of the expected signals
            # arrive within ``boot_wait`` the call still returns (timeout IS
            # the fallback — see issue #2265), preserving the original
            # "waited N seconds then moved on" behavior.
            assert self._serial is not None
            try:
                matched, poll_drained_lines = await wait_for_signal(
                    self._serial,
                    effective_signals,
                    timeout=boot_wait,
                    on_line=_mirror if self.verbose else None,
                )
            except KeyboardInterrupt:
                notify_main_thread()
                raise

            self._check_interrupt()

            if self.verbose:
                if matched is not None:
                    print(
                        f"✅ [RPC] Boot signal {matched!r} detected after "
                        f"draining {len(poll_drained_lines)} line(s); "
                        f"skipping fixed sleep."
                    )
                else:
                    print(
                        f"⏳ [RPC] No boot signal in {boot_wait}s "
                        f"(drained {len(poll_drained_lines)} line(s)); "
                        f"falling back to timeout behavior."
                    )

        if drain_boot:
            if self.verbose:
                print(f"📥 [RPC] Draining boot output...")
            lines_drained = await self.drain_boot_output(verbose=self.verbose)
            if self.verbose:
                total = lines_drained + len(poll_drained_lines)
                print(
                    f"📥 [RPC] Drained {lines_drained} boot lines "
                    f"(plus {len(poll_drained_lines)} during spin-poll = {total})"
                )

    async def close(self) -> None:
        """Close serial connection (async)."""
        self._connected = False
        if self._serial is not None:
            await self._serial.close()
            self._serial = None
            await asyncio.sleep(0)  # Yield control

    async def drain_boot_output(
        self, max_lines: int = 100, verbose: bool = False
    ) -> int:
        """Drain pending boot output from serial buffer (async).

        Args:
            max_lines: Maximum lines to drain
            verbose: Print drained lines

        Returns:
            Number of lines drained
        """
        if not self.is_connected:
            raise RpcError("Not connected")

        assert self._serial is not None

        lines_drained = 0

        # Poll for boot output with short timeout (wait for buffer to drain)
        # Use read_lines with a short timeout to drain existing buffer
        try:
            async for line in self._serial.read_lines(timeout=1.0):
                self._check_interrupt()
                lines_drained += 1

                if verbose:
                    if lines_drained <= 3:
                        print(f"  [boot] {line[:80]}")
                    elif lines_drained == 4:
                        print("  [boot] ... (draining boot output)")

                if lines_drained >= max_lines:
                    break

        except KeyboardInterrupt:
            notify_main_thread()
            raise
        except Exception:
            pass  # Timeout is normal when buffer is drained

        return lines_drained

    async def send(
        self,
        function: str,
        args: list[Any] | None = None,
        timeout: float | None = None,
        retries: int = 1,
    ) -> RpcResponse:
        """Send JSON-RPC command and wait for response with mandatory ID correlation (async).

        Args:
            function: RPC function name
            args: Function arguments as list (default: [])
            timeout: Override default timeout for this call
            retries: Number of retry attempts (default: 1 = no retries)

        Returns:
            RpcResponse with parsed data and matching request ID (ID always present)

        Raises:
            RpcTimeoutError: If no response within timeout
            RpcError: If not connected
        """
        if not self.is_connected:
            raise RpcError("Not connected")

        assert self._serial is not None

        # Generate unique request ID for JSON-RPC 2.0 correlation (mandatory)
        request_id = self._next_id
        self._next_id = (self._next_id + 1) & 0xFFFFFFFF  # Wrap at uint32 max

        # Wrap args in array to pass as single Json parameter to RPC functions
        # RPC functions expect signature: (const fl::json& args)
        # So we need to pass the entire args as one parameter
        wrapped_args: list[Any] = [args] if args is not None else [{}]
        cmd: dict[str, Any] = {
            "method": function,
            "params": wrapped_args,
            "id": request_id,  # ID is mandatory
        }
        cmd_str = json.dumps(cmd, separators=(",", ":"))

        if self.verbose:
            print(f"📤 [RPC] Sending: {cmd_str}")

        effective_timeout = timeout if timeout is not None else self.timeout
        attempt_timeout = effective_timeout / max(retries, 1)

        last_error: Exception | None = None

        for _attempt in range(retries):
            try:
                # Write command via serial interface
                await self._serial.write(cmd_str + "\n")

                # Await response with ID matching (mandatory)
                response = await self._wait_for_response(
                    attempt_timeout, expected_id=request_id
                )
                return response

            except RpcTimeoutError as e:
                last_error = e
                continue

        raise last_error or RpcTimeoutError(
            f"No response after {retries} attempts (total timeout: {effective_timeout}s)"
        )

    async def send_and_match(
        self,
        function: str,
        args: list[Any] | None = None,
        match_key: str | None = None,
        timeout: float | None = None,
        retries: int = 3,
    ) -> RpcResponse:
        """Send RPC and wait for response containing specific key with mandatory ID correlation (async).

        Useful when device may send multiple response types and you need
        to match a specific one.

        Args:
            function: RPC function name
            args: Function arguments
            match_key: Key that must be present in response (e.g., "connected")
            timeout: Override default timeout
            retries: Number of retry attempts

        Returns:
            RpcResponse containing the match_key and matching request ID (always includes ID)

        Raises:
            RpcTimeoutError: If matching response not received
        """
        if not self.is_connected:
            raise RpcError("Not connected")

        assert self._serial is not None

        # Generate unique request ID for JSON-RPC 2.0 correlation (mandatory)
        request_id = self._next_id
        self._next_id = (self._next_id + 1) & 0xFFFFFFFF  # Wrap at uint32 max

        # Wrap args in array to pass as single Json parameter to RPC functions
        # RPC functions expect signature: (const fl::json& args)
        # So we need to pass the entire args as one parameter
        wrapped_args: list[Any] = [args] if args is not None else [{}]
        cmd: dict[str, Any] = {
            "method": function,
            "params": wrapped_args,
            "id": request_id,  # ID is mandatory
        }
        cmd_str = json.dumps(cmd, separators=(",", ":"))

        if self.verbose:
            print(f"📤 [RPC] Sending: {cmd_str}")

        effective_timeout = timeout if timeout is not None else self.timeout
        attempt_timeout = effective_timeout / max(retries, 1)

        for _attempt in range(retries):
            self._check_interrupt()  # Check before each attempt

            # Write command via serial interface
            await self._serial.write(cmd_str + "\n")

            start = time.time()
            try:
                async for line in self._serial.read_lines(timeout=attempt_timeout):
                    self._check_interrupt()

                    # Show all serial output in verbose mode (helps debug crashes)
                    if self.verbose and not line.startswith(self.RESPONSE_PREFIX):
                        print(f"  [serial] {line[:120]}")

                    # Feed non-REMOTE lines to crash decoder.
                    if self._crash_decoder is not None and not line.startswith(
                        self.RESPONSE_PREFIX
                    ):
                        decoded = self._crash_decoder.process_line(line)
                        if decoded:
                            for dl in decoded:
                                print(dl)
                            raise RpcCrashError(
                                "Device crashed during RPC operation",
                                decoded_lines=decoded,
                            )

                    if line.startswith(self.RESPONSE_PREFIX):
                        json_str = line[len(self.RESPONSE_PREFIX) :]
                        try:
                            data = json.loads(json_str)

                            # Check if response ID matches request ID (JSON-RPC 2.0 correlation)
                            # ID is mandatory - must match
                            response_id = data.get("id")
                            if response_id is None or response_id != request_id:
                                # Skip responses without ID or non-matching IDs
                                continue

                            # Check if match_key is present
                            # For JSON-RPC 2.0 responses, check inside "result" field
                            match_found = False
                            if match_key is None:
                                match_found = True
                            elif match_key in data:
                                match_found = True
                            elif (
                                "result" in data
                                and isinstance(data["result"], dict)
                                and match_key in data["result"]
                            ):
                                match_found = True

                            if match_found:
                                # Extract result/error field for JSON-RPC 2.0 responses
                                # Priority: result > error > empty dict (no protocol fields)
                                if "result" in data:
                                    response_data = data["result"]
                                elif "error" in data:
                                    response_data = data["error"]
                                else:
                                    # No result or error - empty response
                                    response_data = {}

                                # Determine success: void functions return null, treat as success
                                if "success" in data:
                                    success = data.get("success", True)  # type: ignore[assignment]
                                elif isinstance(response_data, dict):
                                    success = response_data.get("success", True)  # type: ignore[assignment]
                                else:
                                    # Void functions return null - treat as successful execution
                                    success = True

                                # Ensure response_data is always a dict for consistent API
                                if response_data is None or not isinstance(
                                    response_data, dict
                                ):
                                    response_data = {}

                                return RpcResponse(
                                    success=success,  # type: ignore[arg-type]
                                    data=response_data,  # type: ignore[arg-type]
                                    raw_line=line,
                                    _id=response_id,  # ID is always present
                                )
                        except json.JSONDecodeError:
                            continue

                    # Check timeout
                    if time.time() - start >= attempt_timeout:
                        break

            except KeyboardInterrupt:
                notify_main_thread()
                raise
            except Exception:
                continue

        raise RpcTimeoutError(
            f"No response with key '{match_key}' after {retries} attempts"
        )

    async def _wait_for_response(self, timeout: float, expected_id: int) -> RpcResponse:
        """Wait for and parse JSON-RPC response with mandatory ID matching (async).

        Args:
            timeout: Maximum time to wait in seconds
            expected_id: Expected request ID for JSON-RPC 2.0 correlation (mandatory)

        Returns:
            RpcResponse with parsed data and matching request ID (ID always present)

        Raises:
            RpcTimeoutError: If no matching response within timeout
        """
        assert self._serial is not None

        start = time.time()

        try:
            async for line in self._serial.read_lines(timeout=timeout):
                self._check_interrupt()

                if not line.startswith(self.RESPONSE_PREFIX):
                    if self.verbose:
                        print(f"  [async-serial] {line[:200]}")
                    # Feed line to crash decoder if available.
                    if self._crash_decoder is not None:
                        decoded = self._crash_decoder.process_line(line)
                        if decoded:
                            for dl in decoded:
                                print(dl)
                            raise RpcCrashError(
                                "Device crashed during RPC operation",
                                decoded_lines=decoded,
                            )
                    # Check timeout for non-REMOTE lines
                    if time.time() - start >= timeout:
                        break
                    continue

                # Line starts with RESPONSE_PREFIX - parse as JSON-RPC response
                json_str = line[len(self.RESPONSE_PREFIX) :]
                try:
                    data = json.loads(json_str)

                    # Check if response ID matches expected ID (JSON-RPC 2.0 correlation)
                    # ID is mandatory - must match exactly
                    response_id = data.get("id")
                    if response_id is None:
                        # Responses without ID are invalid - skip them
                        continue
                    if response_id != expected_id:
                        # Skip responses that don't match our expected request ID
                        continue

                    # JSON-RPC 2.0: "error" and "result" are mutually exclusive
                    # If "error" is present, raise RpcError immediately
                    if "error" in data:
                        error_obj = data["error"]

                        # Extract error details (code and message are standard fields)
                        if isinstance(error_obj, dict):
                            error_code = error_obj.get("code")  # type: ignore[assignment]
                            error_message = error_obj.get(  # type: ignore[assignment]
                                "message", "Unknown error"
                            )
                            error_data = error_obj.get("data")  # type: ignore[assignment]

                            # Build detailed error message
                            msg = f"RPC Error {error_code}: {error_message}"
                            if error_data:
                                import json as json_module

                                msg += f" (data: {json_module.dumps(error_data)})"

                            raise RpcError(msg)
                        else:
                            # Malformed error object (should be dict)
                            raise RpcError(f"Malformed error object: {error_obj}")

                    # Extract result field for JSON-RPC 2.0 responses
                    if "result" in data:
                        response_data = data["result"]
                    else:
                        # No result or error - empty response
                        response_data = {}

                    # Check if this is an ACK for async function (skip and continue waiting)
                    if (
                        isinstance(response_data, dict)
                        and "acknowledged" in response_data
                        and response_data["acknowledged"] is True
                    ):
                        # Async RPCs (Phase 3: runSingleTest, runParallelTest,
                        # all testCoroutine*, the net/ble test runners) can run
                        # for tens of seconds — the runSingleTest path through
                        # 100 LEDs × 4 patterns takes ~80s on Teensy 4. Folding
                        # both the ACK-wait and the result-wait into the single
                        # caller-supplied `timeout` caused
                        # `RpcTimeoutError('No response with ID N within 60s')`
                        # for tests that genuinely needed 70-90s of post-ACK
                        # time. Restart the wait with a budget large enough to
                        # cover any reasonable test: the caller's own timeout
                        # if they bumped it above the async floor, else
                        # `ASYNC_POST_ACK_TIMEOUT` (10 min). See FastLED#3060.
                        post_ack_timeout = max(timeout, self.ASYNC_POST_ACK_TIMEOUT)
                        if self.verbose:
                            print(
                                f"[RPC] ACK received for request {expected_id}, "
                                f"restarting wait for final response "
                                f"(post-ACK timeout={post_ack_timeout}s)..."
                            )
                        return await self._wait_for_response(
                            post_ack_timeout, expected_id
                        )

                    # Determine success: void functions return null, treat as success
                    success: bool  # Explicit type declaration
                    if "success" in data:
                        success = bool(data.get("success", True))
                    elif isinstance(response_data, dict):
                        success = bool(response_data.get("success", True))  # type: ignore[arg-type]
                    else:
                        # Void functions return null - treat as successful execution
                        success = True

                    # Ensure response_data is always a dict for consistent API
                    final_response_data: dict[str, Any]
                    if response_data is None or not isinstance(response_data, dict):
                        final_response_data = {}
                    else:
                        final_response_data = response_data  # type: ignore[assignment]

                    return RpcResponse(
                        success=success,
                        data=final_response_data,
                        raw_line=line,
                        _id=response_id,  # ID is always present
                    )
                except json.JSONDecodeError:
                    continue

        except KeyboardInterrupt:
            notify_main_thread()
            raise
        except RpcCrashError:
            raise  # Re-raise crash errors without masking them
        except Exception as e:
            if self.verbose:
                print(f"⚠️  [RPC] Exception during read_lines: {e}")
            pass  # Timeout or other error

        # Flush any partial crash data before raising timeout.
        if self._crash_decoder is not None:
            flushed = self._crash_decoder.flush()
            if flushed:
                for fl in flushed:
                    print(fl)
                raise RpcCrashError(
                    "Device crashed during RPC operation (detected at timeout)",
                    decoded_lines=flushed,
                )

        if self.verbose:
            print(
                f"❌ [RPC] Timeout: No response with ID {expected_id} within {timeout}s"
            )
        raise RpcTimeoutError(f"No response with ID {expected_id} within {timeout}s")

    def _check_interrupt(self) -> None:
        """Check for KeyboardInterrupt signal."""
        if is_interrupted():
            raise KeyboardInterrupt()

    async def __aenter__(self) -> "RpcClient":
        """Async context manager entry - connects to device."""
        await self.connect()
        return self

    async def __aexit__(
        self, _exc_type: type | None, _exc_val: Exception | None, _exc_tb: object
    ) -> None:
        """Async context manager exit - closes connection."""
        await self.close()
