"""WebSocket connection management with Redis-backed session tracking.

This module manages WebSocket connection state using Redis for persistence
and cross-instance coordination. It tracks:
- Individual connection sessions (sid -> user_id, game_id, timestamps)
- User to connection mapping (user_id -> sid)
- Game participants (game_id -> set of sids)

Key Patterns:
    conn:{sid} - Hash with connection details (user_id, game_id, connected_at, last_seen)
    user_conn:{user_id} - String with active sid
    game_conns:{game_id} - Set of sids for game participants

Lifecycle:
    1. on_connect: Register connection, map user to sid
    2. on_join_game: Add sid to game's connection set
    3. on_heartbeat: Update last_seen timestamp
    4. on_disconnect: Remove from all tracking structures

Example:
    manager = ConnectionManager()

    # On WebSocket connect
    await manager.register_connection(sid, user_id)

    # On joining a game
    await manager.join_game(sid, game_id)

    # Check if opponent is online
    if await manager.is_user_online(opponent_id):
        # Send real-time update

    # On disconnect
    await manager.unregister_connection(sid)
"""

import logging
from collections.abc import AsyncIterator, Callable
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING
from uuid import UUID

from app.db.redis import get_redis

if TYPE_CHECKING:
    from redis.asyncio import Redis

logger = logging.getLogger(__name__)

# Type alias for redis factory - a callable that returns an async context manager
RedisFactory = Callable[[], AsyncIterator["Redis"]]

# Redis key patterns
CONN_PREFIX = "conn:"
USER_CONN_PREFIX = "user_conn:"
GAME_CONNS_PREFIX = "game_conns:"
SPECTATORS_PREFIX = "spectators:"

# Connection TTL (auto-expire stale connections)
DEFAULT_CONN_TTL_SECONDS = 3600  # 1 hour
HEARTBEAT_INTERVAL_SECONDS = 30


@dataclass
class ConnectionInfo:
    """Information about a WebSocket connection.

    Attributes:
        sid: Socket.IO session ID.
        user_id: Authenticated user's UUID.
        game_id: Current game ID if in a game, None otherwise.
        connected_at: When the connection was established.
        last_seen: Last heartbeat or activity timestamp.
    """

    sid: str
    user_id: str
    game_id: str | None
    connected_at: datetime
    last_seen: datetime

    def is_stale(self, threshold_seconds: int = HEARTBEAT_INTERVAL_SECONDS * 3) -> bool:
        """Check if connection is stale (no recent activity).

        Args:
            threshold_seconds: Seconds since last_seen to consider stale.

        Returns:
            True if connection hasn't been seen recently.
        """
        elapsed = (datetime.now(UTC) - self.last_seen).total_seconds()
        return elapsed > threshold_seconds


class ConnectionManager:
    """Manages WebSocket connections with Redis-backed session tracking.

    Provides methods to:
    - Register/unregister connections
    - Track which game a connection is participating in
    - Check user online status
    - Get all connections for a game
    - Handle heartbeats and stale connection detection

    Attributes:
        conn_ttl_seconds: TTL for connection records in Redis.
    """

    def __init__(
        self,
        conn_ttl_seconds: int = DEFAULT_CONN_TTL_SECONDS,
        redis_factory: RedisFactory | None = None,
    ) -> None:
        """Initialize the ConnectionManager.

        Args:
            conn_ttl_seconds: How long to keep connection records in Redis.
            redis_factory: Optional factory for Redis connections. If not provided,
                uses the default get_redis from app.db.redis. Useful for testing.
        """
        self.conn_ttl_seconds = conn_ttl_seconds
        self._get_redis = redis_factory if redis_factory is not None else get_redis

    def _conn_key(self, sid: str) -> str:
        """Generate Redis key for a connection."""
        return f"{CONN_PREFIX}{sid}"

    def _user_conn_key(self, user_id: str) -> str:
        """Generate Redis key for user-to-connection mapping."""
        return f"{USER_CONN_PREFIX}{user_id}"

    def _game_conns_key(self, game_id: str) -> str:
        """Generate Redis key for game connection set."""
        return f"{GAME_CONNS_PREFIX}{game_id}"

    def _spectators_key(self, game_id: str) -> str:
        """Generate Redis key for game spectators set."""
        return f"{SPECTATORS_PREFIX}{game_id}"

    # =========================================================================
    # Connection Lifecycle
    # =========================================================================

    async def register_connection(
        self,
        sid: str,
        user_id: str | UUID,
    ) -> None:
        """Register a new WebSocket connection.

        Creates connection record in Redis and maps user to this connection.
        If user had a previous connection, it will be replaced.

        Args:
            sid: Socket.IO session ID.
            user_id: Authenticated user's ID (UUID or string).

        Example:
            await manager.register_connection("abc123", user.id)
        """
        user_id_str = str(user_id)
        now = datetime.now(UTC).isoformat()

        async with self._get_redis() as redis:
            # Check for existing connection and clean it up
            old_sid = await redis.get(self._user_conn_key(user_id_str))
            if old_sid and old_sid != sid:
                logger.info(f"Replacing old connection {old_sid} for user {user_id_str}")
                await self._cleanup_connection(old_sid)

            # Create connection record (hash)
            conn_key = self._conn_key(sid)
            await redis.hset(
                conn_key,
                mapping={
                    "user_id": user_id_str,
                    "game_id": "",
                    "connected_at": now,
                    "last_seen": now,
                },
            )
            await redis.expire(conn_key, self.conn_ttl_seconds)

            # Map user to this connection
            user_conn_key = self._user_conn_key(user_id_str)
            await redis.set(user_conn_key, sid)
            await redis.expire(user_conn_key, self.conn_ttl_seconds)

        logger.debug(f"Registered connection: sid={sid}, user_id={user_id_str}")

    async def unregister_connection(self, sid: str) -> ConnectionInfo | None:
        """Unregister a WebSocket connection and clean up all related data.

        Removes connection from:
        - Connection record
        - User-to-connection mapping
        - Any game connection sets

        Args:
            sid: Socket.IO session ID.

        Returns:
            ConnectionInfo of the removed connection, or None if not found.

        Example:
            info = await manager.unregister_connection("abc123")
            if info and info.game_id:
                # Notify opponent of disconnect
        """
        # Get connection info before cleanup
        info = await self.get_connection(sid)
        if info is None:
            logger.debug(f"Connection not found for unregister: {sid}")
            return None

        await self._cleanup_connection(sid)
        logger.debug(f"Unregistered connection: sid={sid}, user_id={info.user_id}")
        return info

    async def _cleanup_connection(self, sid: str) -> None:
        """Internal cleanup of a connection's Redis data.

        Args:
            sid: Socket.IO session ID.
        """
        async with self._get_redis() as redis:
            conn_key = self._conn_key(sid)

            # Get connection data for cleanup
            data = await redis.hgetall(conn_key)
            if not data:
                return

            user_id = data.get("user_id", "")
            game_id = data.get("game_id", "")

            # Warn if user_id is missing - indicates corrupted data
            if not user_id:
                logger.warning(f"Connection {sid} has no user_id - data may be corrupted")

            # Check if spectating (game_id format: "spectating:{actual_game_id}")
            if game_id and game_id.startswith("spectating:"):
                actual_game_id = game_id[len("spectating:") :]
                spectators_key = self._spectators_key(actual_game_id)
                await redis.srem(spectators_key, sid)
            elif game_id:
                # Remove from game connection set if in a game as player
                game_conns_key = self._game_conns_key(game_id)
                await redis.srem(game_conns_key, sid)

            # Remove user-to-connection mapping if it points to this sid
            if user_id:
                user_conn_key = self._user_conn_key(user_id)
                current_sid = await redis.get(user_conn_key)
                if current_sid == sid:
                    await redis.delete(user_conn_key)

            # Delete connection record
            await redis.delete(conn_key)

    # =========================================================================
    # Game Association
    # =========================================================================

    async def join_game(self, sid: str, game_id: str) -> bool:
        """Associate a connection with a game.

        Updates the connection's game_id and adds it to the game's connection set.
        If connection was in a different game, leaves that game first.

        Args:
            sid: Socket.IO session ID.
            game_id: Game to join.

        Returns:
            True if successful, False if connection not found.

        Example:
            await manager.join_game("abc123", "game-456")
        """
        async with self._get_redis() as redis:
            conn_key = self._conn_key(sid)

            # Check connection exists
            exists = await redis.exists(conn_key)
            if not exists:
                logger.warning(f"Cannot join game: connection not found {sid}")
                return False

            # Leave current game if in one
            current_game = await redis.hget(conn_key, "game_id")
            if current_game and current_game != game_id:
                old_game_conns_key = self._game_conns_key(current_game)
                await redis.srem(old_game_conns_key, sid)
                logger.debug(f"Connection {sid} left game {current_game}")

            # Update connection's game_id
            await redis.hset(conn_key, "game_id", game_id)

            # Add to game's connection set
            game_conns_key = self._game_conns_key(game_id)
            await redis.sadd(game_conns_key, sid)

            # Set TTL on game conns set (cleanup if game becomes inactive)
            await redis.expire(game_conns_key, self.conn_ttl_seconds)

        logger.debug(f"Connection {sid} joined game {game_id}")
        return True

    async def leave_game(self, sid: str) -> str | None:
        """Remove a connection from its current game.

        Args:
            sid: Socket.IO session ID.

        Returns:
            The game_id that was left, or None if not in a game.

        Example:
            game_id = await manager.leave_game("abc123")
        """
        async with self._get_redis() as redis:
            conn_key = self._conn_key(sid)

            # Get current game
            game_id = await redis.hget(conn_key, "game_id")
            if not game_id:
                return None

            # Remove from game's connection set
            game_conns_key = self._game_conns_key(game_id)
            await redis.srem(game_conns_key, sid)

            # Clear game_id on connection
            await redis.hset(conn_key, "game_id", "")

        logger.debug(f"Connection {sid} left game {game_id}")
        return game_id

    # =========================================================================
    # Heartbeat / Activity
    # =========================================================================

    async def update_heartbeat(self, sid: str) -> bool:
        """Update the last_seen timestamp for a connection.

        Should be called on each heartbeat message from the client.
        Also refreshes TTL on connection records.

        Args:
            sid: Socket.IO session ID.

        Returns:
            True if successful, False if connection not found.

        Example:
            await manager.update_heartbeat("abc123")
        """
        now = datetime.now(UTC).isoformat()

        async with self._get_redis() as redis:
            conn_key = self._conn_key(sid)

            # Check exists and update atomically
            exists = await redis.exists(conn_key)
            if not exists:
                return False

            # Update last_seen
            await redis.hset(conn_key, "last_seen", now)

            # Refresh TTL
            await redis.expire(conn_key, self.conn_ttl_seconds)

            # Also refresh user mapping TTL
            user_id = await redis.hget(conn_key, "user_id")
            if user_id:
                user_conn_key = self._user_conn_key(user_id)
                await redis.expire(user_conn_key, self.conn_ttl_seconds)

        return True

    # =========================================================================
    # Query Methods
    # =========================================================================

    async def get_connection(self, sid: str) -> ConnectionInfo | None:
        """Get connection info by session ID.

        Args:
            sid: Socket.IO session ID.

        Returns:
            ConnectionInfo if found, None otherwise.

        Example:
            info = await manager.get_connection("abc123")
            if info:
                print(f"User {info.user_id} connected at {info.connected_at}")
        """
        async with self._get_redis() as redis:
            conn_key = self._conn_key(sid)
            data = await redis.hgetall(conn_key)

            if not data:
                return None

            # Validate required fields - if missing, data is corrupted
            user_id = data.get("user_id")
            connected_at = data.get("connected_at")
            last_seen = data.get("last_seen")

            if not user_id or not connected_at or not last_seen:
                logger.warning(
                    f"Corrupted connection record for {sid}: "
                    f"missing required fields (user_id={bool(user_id)}, "
                    f"connected_at={bool(connected_at)}, last_seen={bool(last_seen)})"
                )
                return None

            return ConnectionInfo(
                sid=sid,
                user_id=user_id,
                game_id=data.get("game_id") or None,
                connected_at=datetime.fromisoformat(connected_at),
                last_seen=datetime.fromisoformat(last_seen),
            )

    async def get_user_connection(self, user_id: str | UUID) -> ConnectionInfo | None:
        """Get connection info for a user.

        Args:
            user_id: User's UUID or string ID.

        Returns:
            ConnectionInfo if user is connected, None otherwise.

        Example:
            info = await manager.get_user_connection(opponent_id)
            if info:
                # User is online
        """
        user_id_str = str(user_id)

        async with self._get_redis() as redis:
            user_conn_key = self._user_conn_key(user_id_str)
            sid = await redis.get(user_conn_key)

            if not sid:
                return None

        return await self.get_connection(sid)

    async def is_user_online(self, user_id: str | UUID) -> bool:
        """Check if a user has an active connection.

        Args:
            user_id: User's UUID or string ID.

        Returns:
            True if user is connected, False otherwise.

        Example:
            if await manager.is_user_online(opponent_id):
                # Send real-time notification
        """
        info = await self.get_user_connection(user_id)
        if info is None:
            return False

        # Check if connection is stale
        if info.is_stale():
            logger.debug(f"User {user_id} connection is stale")
            return False

        return True

    async def get_game_connections(self, game_id: str) -> list[ConnectionInfo]:
        """Get all connections for a game.

        Args:
            game_id: Game ID.

        Returns:
            List of ConnectionInfo for all connected participants.

        Example:
            connections = await manager.get_game_connections("game-456")
            for conn in connections:
                # Send state update to each participant
        """
        async with self._get_redis() as redis:
            game_conns_key = self._game_conns_key(game_id)
            sids = await redis.smembers(game_conns_key)

        connections = []
        for sid in sids:
            info = await self.get_connection(sid)
            if info is not None:
                connections.append(info)

        return connections

    async def get_game_user_sids(self, game_id: str) -> dict[str, str]:
        """Get mapping of user_id to sid for a game.

        Useful for targeted message delivery to specific players.

        Args:
            game_id: Game ID.

        Returns:
            Dict mapping user_id to their sid.

        Example:
            user_sids = await manager.get_game_user_sids("game-456")
            player1_sid = user_sids.get(player1_id)
        """
        connections = await self.get_game_connections(game_id)
        return {conn.user_id: conn.sid for conn in connections}

    async def get_opponent_sid(
        self,
        game_id: str,
        current_user_id: str | UUID,
    ) -> str | None:
        """Get the sid of the opponent in a 2-player game.

        Args:
            game_id: Game ID.
            current_user_id: The current user's ID.

        Returns:
            Opponent's sid if found and connected, None otherwise.

        Example:
            opponent_sid = await manager.get_opponent_sid("game-456", user_id)
            if opponent_sid:
                # Send message to opponent
        """
        current_user_str = str(current_user_id)
        connections = await self.get_game_connections(game_id)

        for conn in connections:
            if conn.user_id != current_user_str:
                return conn.sid

        return None

    async def get_user_active_game(self, user_id: str | UUID) -> str | None:
        """Get the game ID of a user's active game from their connection.

        This checks the ConnectionManager's tracking, NOT the database.
        Used during reconnection to find if the user was in a game.

        Args:
            user_id: User's UUID or string ID.

        Returns:
            Game ID if user was in a game, None otherwise.

        Example:
            game_id = await manager.get_user_active_game(user_id)
            if game_id:
                # User was in a game, auto-rejoin
        """
        conn_info = await self.get_user_connection(user_id)
        if conn_info is None:
            return None
        return conn_info.game_id

    # =========================================================================
    # Maintenance
    # =========================================================================

    async def cleanup_stale_connections(
        self,
        threshold_seconds: int = HEARTBEAT_INTERVAL_SECONDS * 3,
    ) -> int:
        """Clean up stale connections that haven't sent heartbeats.

        Should be called periodically by a background task.

        Args:
            threshold_seconds: Seconds since last_seen to consider stale.

        Returns:
            Number of connections cleaned up.

        Example:
            # In background task
            count = await manager.cleanup_stale_connections()
            logger.info(f"Cleaned up {count} stale connections")
        """
        count = 0
        async with self._get_redis() as redis:
            # Scan for all connection keys
            async for key in redis.scan_iter(match=f"{CONN_PREFIX}*"):
                sid = key[len(CONN_PREFIX) :]
                info = await self.get_connection(sid)

                if info and info.is_stale(threshold_seconds):
                    await self._cleanup_connection(sid)
                    count += 1
                    logger.debug(f"Cleaned up stale connection: {sid}")

        if count > 0:
            logger.info(f"Cleaned up {count} stale connections")

        return count

    async def get_connection_count(self) -> int:
        """Get the total number of active connections.

        Useful for monitoring and admin dashboards.

        Returns:
            Number of active connections.
        """
        count = 0
        async with self._get_redis() as redis:
            async for _ in redis.scan_iter(match=f"{CONN_PREFIX}*"):
                count += 1
        return count

    async def get_game_connection_count(self, game_id: str) -> int:
        """Get the number of connections for a specific game.

        Args:
            game_id: Game ID.

        Returns:
            Number of connected participants.
        """
        async with self._get_redis() as redis:
            game_conns_key = self._game_conns_key(game_id)
            return await redis.scard(game_conns_key)

    # =========================================================================
    # Spectator Management
    # =========================================================================

    async def register_spectator(
        self,
        sid: str,
        user_id: str | UUID,
        game_id: str,
    ) -> bool:
        """Register a connection as a spectator for a game.

        Adds the sid to the spectators set for the game and updates
        the connection record with spectator status.

        Args:
            sid: Socket.IO session ID.
            user_id: User's ID (for logging).
            game_id: Game to spectate.

        Returns:
            True if successful, False if connection not found.

        Example:
            await manager.register_spectator("abc123", user_id, "game-456")
        """
        user_id_str = str(user_id)

        async with self._get_redis() as redis:
            conn_key = self._conn_key(sid)

            # Check connection exists
            exists = await redis.exists(conn_key)
            if not exists:
                logger.warning(f"Cannot register spectator: connection not found {sid}")
                return False

            # Update connection's game_id to indicate spectating
            # We use a special format to distinguish from playing
            await redis.hset(conn_key, "game_id", f"spectating:{game_id}")

            # Add to spectators set
            spectators_key = self._spectators_key(game_id)
            await redis.sadd(spectators_key, sid)

            # Set TTL on spectators set
            await redis.expire(spectators_key, self.conn_ttl_seconds)

        logger.debug(f"User {user_id_str} ({sid}) spectating game {game_id}")
        return True

    async def unregister_spectator(self, sid: str, game_id: str) -> bool:
        """Remove a connection from a game's spectator list.

        Args:
            sid: Socket.IO session ID.
            game_id: Game being spectated.

        Returns:
            True if removed, False if not in spectator list.

        Example:
            await manager.unregister_spectator("abc123", "game-456")
        """
        async with self._get_redis() as redis:
            spectators_key = self._spectators_key(game_id)

            # Remove from spectators set
            removed = await redis.srem(spectators_key, sid)

            # Clear game_id on connection if it was spectating this game
            conn_key = self._conn_key(sid)
            current_game = await redis.hget(conn_key, "game_id")
            if current_game == f"spectating:{game_id}":
                await redis.hset(conn_key, "game_id", "")

        if removed:
            logger.debug(f"Connection {sid} stopped spectating game {game_id}")

        return bool(removed)

    async def get_spectator_count(self, game_id: str) -> int:
        """Get the number of spectators for a game.

        Args:
            game_id: Game ID.

        Returns:
            Number of spectators.

        Example:
            count = await manager.get_spectator_count("game-456")
        """
        async with self._get_redis() as redis:
            spectators_key = self._spectators_key(game_id)
            return await redis.scard(spectators_key)

    async def get_game_spectators(self, game_id: str) -> list[str]:
        """Get all spectator sids for a game.

        Args:
            game_id: Game ID.

        Returns:
            List of spectator sids.

        Example:
            sids = await manager.get_game_spectators("game-456")
        """
        async with self._get_redis() as redis:
            spectators_key = self._spectators_key(game_id)
            sids = await redis.smembers(spectators_key)
            return list(sids)

    async def is_spectating(self, sid: str, game_id: str) -> bool:
        """Check if a connection is spectating a specific game.

        Args:
            sid: Socket.IO session ID.
            game_id: Game ID.

        Returns:
            True if spectating, False otherwise.
        """
        async with self._get_redis() as redis:
            spectators_key = self._spectators_key(game_id)
            return await redis.sismember(spectators_key, sid)


# Global singleton instance
connection_manager = ConnectionManager()
