"""Per-client WebSocket connection state and inbound dispatch."""

from __future__ import annotations

import asyncio
import logging
import time
from typing import Any, Dict, Optional

from aiohttp import web, WSMsgType

from . import protocol as proto

logger = logging.getLogger(__name__)


class ClientConnection:
    """One authenticated WebSocket client (single-user app: at most a few)."""

    def __init__(self, ws: "web.WebSocketResponse", conn_id: str):
        self.ws = ws
        self.conn_id = conn_id
        self.authenticated = False
        self.protocol_version: Optional[int] = None
        self.connected_at = time.time()
        # Serialized outbound writes — aiohttp WS frames are not safe to
        # compose concurrently from multiple tasks.
        self._send_lock = asyncio.Lock()
        self.closed = False

    async def send_json(self, payload: Dict[str, Any]) -> bool:
        """Thread-safe send; returns False when the socket is going away."""
        if self.closed:
            return False
        try:
            async with self._send_lock:
                await self.ws.send_str(proto.encode_message(payload))
            return True
        except (ConnectionError, RuntimeError, asyncio.CancelledError):
            self.closed = True
            return False

    # ------------------------------------------------------------------
    async def read_loop(self, server: Any) -> None:
        """Receive/dispatch loop; exits on close, error, or auth timeout."""
        try:
            async for msg in self.ws:
                if msg.type == WSMsgType.TEXT:
                    if len(msg.data) > proto.MAX_WS_MESSAGE_BYTES:
                        await self.send_json(proto.error_event(
                            proto.ERR_TOO_LARGE,
                            "WebSocket message exceeds server limit"))
                        continue
                    await server.handle_client_message(self, msg.data)
                elif msg.type == WSMsgType.BINARY:
                    await self.send_json(proto.error_event(
                        proto.ERR_BAD_REQUEST,
                        "Binary frames are not part of the Pheby protocol"))
                elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING,
                                  WSMsgType.CLOSED):
                    break
                elif msg.type == WSMsgType.ERROR:
                    logger.warning("[pheby] ws error on %s: %s",
                                   self.conn_id, msg.data)
                    break
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.warning("[pheby] client %s read loop crashed",
                           self.conn_id, exc_info=True)
        finally:
            self.closed = True