"""ACP JSON-RPC 2.0 method router.

Dispatches incoming JSON-RPC requests to the appropriate handler based on
the method prefix: ``session/*`` → SessionHandler, ``device/*`` → DeviceHandler.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Callable, Awaitable

from .device_handler import DeviceHandler
from .session_handler import SessionHandler

logger = logging.getLogger(__name__)

JsonRpcHandler = Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]


class ACPRouter:
    """Routes ACP JSON-RPC 2.0 messages to device and session handlers."""

    def __init__(
        self,
        device_handler: DeviceHandler,
        session_handler: SessionHandler,
    ) -> None:
        self._device = device_handler
        self._session = session_handler
        self._handlers: dict[str, JsonRpcHandler] = {}
        self._register_handlers()

    def _register_handlers(self) -> None:
        self._handlers["device/info"] = self._device.handle_info
        self._handlers["device/screenshot"] = self._device.handle_screenshot
        self._handlers["device/ui-tree"] = self._device.handle_ui_tree
        self._handlers["device/action"] = self._device.handle_action
        self._handlers["device/install"] = self._device.handle_install
        self._handlers["device/shell"] = self._device.handle_shell

        self._handlers["session/prompt"] = self._session.handle_prompt
        self._handlers["session/new"] = self._session.handle_new
        self._handlers["session/load"] = self._session.handle_load
        self._handlers["session/list"] = self._session.handle_list

    async def handle_message(self, raw: str) -> str | None:
        """Parse and dispatch a JSON-RPC 2.0 message, returning the response JSON.

        Returns ``None`` for notifications (no ``id`` field).
        """
        try:
            msg = json.loads(raw)
        except json.JSONDecodeError:
            return self._error_response(None, -32700, "Parse error")

        method = msg.get("method", "")
        msg_id = msg.get("id")

        handler = self._handlers.get(method)
        if handler is None:
            if msg_id is not None:
                return self._error_response(msg_id, -32601, f"Method not found: {method}")
            return None

        try:
            result = await handler(msg)
            if msg_id is not None:
                return json.dumps({
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "result": result,
                })
        except Exception as exc:
            logger.exception("Handler error for %s", method)
            if msg_id is not None:
                return self._error_response(msg_id, -32000, str(exc))
        return None

    @staticmethod
    def _error_response(msg_id: Any, code: int, message: str) -> str:
        return json.dumps({
            "jsonrpc": "2.0",
            "id": msg_id,
            "error": {"code": code, "message": message},
        })
