"""Bridge between the game server and the AI service.

The AI service (``wizard_ai``) connects to ``/internal/ai`` via websocket.
Whenever an AI seat has to act, the server sends it the seat's observation
and waits for an action:

    server -> ai:  {"type": "decide", "request_id": "...", "observation": {...}}
    ai -> server:  {"type": "action", "request_id": "...", "action": {...}}
    ai -> server:  {"type": "error",  "request_id": "...", "message": "..."}

Several AI service instances may connect; requests are distributed round
robin. If no instance is connected, it does not answer in time, or it answers
with an illegal action, a random legal action is played instead so a game
never gets stuck.
"""

from __future__ import annotations

import asyncio
import itertools
import logging
import random
import uuid
from collections.abc import Callable
from typing import Any

from starlette.websockets import WebSocket, WebSocketDisconnect

from wizard_game import Action, Observation, action_from_dict

logger = logging.getLogger(__name__)


class AIConnection:
    def __init__(self, websocket: WebSocket, name: str) -> None:
        self.websocket = websocket
        self.name = name
        self.pending: dict[str, asyncio.Future[dict[str, Any]]] = {}
        self._send_lock = asyncio.Lock()

    async def send(self, message: dict[str, Any]) -> None:
        async with self._send_lock:
            await self.websocket.send_json(message)

    def fail_pending(self) -> None:
        for future in self.pending.values():
            if not future.done():
                future.set_exception(ConnectionError("AI service disconnected"))
        self.pending.clear()


class AIBridge:
    def __init__(
        self,
        timeout: float,
        rng: random.Random | None = None,
        on_change: Callable[[], None] | None = None,
    ) -> None:
        self.timeout = timeout
        # Called whenever an AI service connects or disconnects.
        self._on_change = on_change
        self._connections: list[AIConnection] = []
        self._round_robin = itertools.count()
        self._rng = rng or random.Random()

    @property
    def available(self) -> bool:
        return bool(self._connections)

    async def decide(self, observation: Observation) -> Action:
        """Ask the AI service for an action; falls back to a random legal one."""
        legal = observation.legal_actions
        if not legal:
            raise ValueError("Observation has no legal actions")

        connection = self._pick_connection()
        if connection is None:
            return self._fallback(observation, "no AI service connected", logging.DEBUG)

        request_id = uuid.uuid4().hex
        future: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future()
        connection.pending[request_id] = future
        try:
            await connection.send({
                "type": "decide",
                "request_id": request_id,
                "observation": observation.to_dict(),
            })
            reply = await asyncio.wait_for(future, self.timeout)
        except (TimeoutError, ConnectionError, RuntimeError, WebSocketDisconnect) as exc:
            return self._fallback(observation, f"AI request failed: {exc!r}")
        finally:
            connection.pending.pop(request_id, None)

        if reply.get("type") != "action":
            return self._fallback(observation, f"AI error: {reply.get('message')}")
        try:
            action = action_from_dict(reply.get("action", {}))
        except ValueError as exc:
            return self._fallback(observation, f"malformed AI action: {exc}")
        if action not in legal:
            return self._fallback(observation, f"illegal AI action {action}")
        return action

    async def serve(self, websocket: WebSocket) -> None:
        """Handle one AI service connection until it disconnects."""
        connection = AIConnection(websocket, name=websocket.client.host if websocket.client else "?")
        self._connections.append(connection)
        logger.info("AI service connected from %s (%d connected)", connection.name, len(self._connections))
        self._notify_change()
        try:
            while True:
                message = await websocket.receive_json()
                if not isinstance(message, dict):
                    continue
                if message.get("type") == "hello":
                    logger.info("AI service %s: %s", connection.name, message.get("agent"))
                    continue
                future = connection.pending.get(str(message.get("request_id")))
                if future and not future.done():
                    future.set_result(message)
        except (WebSocketDisconnect, RuntimeError, ValueError):
            pass
        finally:
            self._connections.remove(connection)
            connection.fail_pending()
            logger.info("AI service %s disconnected (%d connected)", connection.name, len(self._connections))
            self._notify_change()

    def _notify_change(self) -> None:
        if self._on_change is not None:
            self._on_change()

    def _pick_connection(self) -> AIConnection | None:
        if not self._connections:
            return None
        return self._connections[next(self._round_robin) % len(self._connections)]

    def _fallback(self, observation: Observation, reason: str, level: int = logging.WARNING) -> Action:
        logger.log(level, "Playing random move for player %d: %s", observation.player, reason)
        return self._rng.choice(observation.legal_actions)