import asyncio
import logging
import time
from typing import Dict, Set, Optional
from fastapi import WebSocket
from dataclasses import dataclass, field

from ..config import settings
from .protocol import (
    MessageType,
    ReadyMessage,
    StatsMessage,
)


@dataclass
class ConnectionInfo:
    websocket: WebSocket
    connected_at_ms: int
    last_active_ms: int = 0
    chunks_received: int = 0
    bytes_received: int = 0
    client_info: dict = field(default_factory=dict)


class ConnectionManager:
    """Manages WebSocket connections with per-connection isolation."""

    def __init__(self, max_connections: int = settings.max_connections):
        self._connections: Dict[str, ConnectionInfo] = {}
        self._active_ids: Set[str] = set()
        self._max_connections = max_connections
        self._logger = logging.getLogger(__name__)
        self._lock = asyncio.Lock()

    async def connect(self, websocket: WebSocket, connection_id: str, **kwargs) -> None:
        """Accept connection and register it."""
        async with self._lock:
            if len(self._connections) >= self._max_connections:
                await websocket.close(code=1013, reason="Too many connections")
                raise RuntimeError(f"Max connections ({self._max_connections}) reached")

            await websocket.accept()
            info = ConnectionInfo(
                websocket=websocket,
                connected_at_ms=int(time.time() * 1000),
                client_info=kwargs,
            )
            self._connections[connection_id] = info
            self._active_ids.add(connection_id)
            self._logger.info(f"Connection established: {connection_id}")

    async def disconnect(self, connection_id: str) -> None:
        """Remove connection and cleanup resources."""
        async with self._lock:
            self._connections.pop(connection_id, None)
            self._active_ids.discard(connection_id)
            self._logger.info(f"Connection closed: {connection_id}")

    def is_active(self, connection_id: str) -> bool:
        return connection_id in self._connections

    def get_connection(self, connection_id: str) -> Optional[ConnectionInfo]:
        return self._connections.get(connection_id)

    def touch(self, connection_id: str, n_chunks: int = 0, n_bytes: int = 0) -> None:
        """Update connection activity timestamps and counters."""
        info = self._connections.get(connection_id)
        if info:
            info.last_active_ms = int(time.time() * 1000)
            info.chunks_received += n_chunks
            info.bytes_received += n_bytes

    async def send_text(self, connection_id: str, message: str) -> bool:
        """Send text message to a specific connection."""
        info = self._connections.get(connection_id)
        if not info:
            return False
        try:
            await info.websocket.send_text(message)
            return True
        except Exception as e:
            self._logger.error(f"Failed to send to {connection_id}: {e}")
            return False

    async def send_json(self, connection_id: str, message: dict) -> bool:
        """Send JSON message to a specific connection."""
        info = self._connections.get(connection_id)
        if not info:
            return False
        try:
            await info.websocket.send_json(message)
            return True
        except Exception as e:
            self._logger.error(f"Failed to send JSON to {connection_id}: {e}")
            return False

    async def broadcast_json(self, message: dict) -> int:
        """Broadcast message to all connections. Returns count of successes."""
        successes = 0
        for cid in list(self._connections.keys()):
            if await self.send_json(cid, message):
                successes += 1
        return successes

    def get_active_connections(self) -> Dict[str, ConnectionInfo]:
        return dict(self._connections)

    def get_connection_count(self) -> int:
        return len(self._connections)

    def get_connection_summary(self) -> dict:
        """Get summary of all connections for monitoring."""
        now = int(time.time() * 1000)
        return {
            cid: {
                "connected_at_ms": info.connected_at_ms,
                "uptime_ms": now - info.connected_at_ms,
                "last_active_ms": info.last_active_ms,
                "chunks_received": info.chunks_received,
                "bytes_received": info.bytes_received,
            }
            for cid, info in self._connections.items()
        }

    async def close_all(self) -> None:
        """Close all connections gracefully."""
        for cid in list(self._connections.keys()):
            try:
                await self._connections[cid].websocket.close(code=1001, reason="Server shutting down")
            except Exception:
                pass
        self._connections.clear()
        self._active_ids.clear()

    async def session_cleanup(self) -> None:
        """Periodic cleanup of inactive connections."""
        now = int(time.time() * 1000)
        for cid in list(self._connections.keys()):
            info = self._connections[cid]
            if now - info.last_active_ms > settings.ws_ping_timeout * 1000 * 2:
                self._logger.warning(f"Cleaning up inactive connection: {cid}")
                await self.disconnect(cid)