"""
WebSocket connection manager for real-time game updates.
Handles WebSocket connections and Redis pub/sub for cross-worker communication.
"""
import uuid
from typing import Dict, Optional, List
from fastapi import WebSocket
import logging

logger = logging.getLogger(__name__)


class ConnectionManager:
    """
    Manages WebSocket connections for real-time game updates.
    Tracks active connections and provides methods for broadcasting messages.
    Uses Redis pub/sub for cross-worker communication.
    """
    def __init__(self):
        # Maps player_id -> List of WebSocket connections (local to this worker only)
        self.active_connections: Dict[int, List[WebSocket]] = {}
        # Maps player_id -> username for debugging
        self.player_usernames: Dict[int, str] = {}
        # Redis manager instance (injected later)
        self.redis_manager = None
    
    def set_redis_manager(self, redis_manager):
        """Inject Redis manager after initialization."""
        self.redis_manager = redis_manager
    
    async def connect(self, websocket: WebSocket, player_id: int, username: str):
        """Accept a new WebSocket connection and track it."""
        await websocket.accept()
        
        if player_id not in self.active_connections:
            self.active_connections[player_id] = []
        
        self.active_connections[player_id].append(websocket)
        self.player_usernames[player_id] = username
        
        # Subscribe to player's personal channel (only if first connection)
        if len(self.active_connections[player_id]) == 1 and self.redis_manager:
            await self.redis_manager.subscribe_to_channels([f"player:{player_id}"])
            await self.redis_manager.mark_player_connected(player_id)
        
        logger.info(f"WebSocket connected: {username} (player_id={player_id}, worker={self.redis_manager.worker_id if self.redis_manager else 'N/A'})")
    
    async def disconnect(self, player_id: int, websocket: WebSocket):
        """Remove a WebSocket connection."""
        if player_id in self.active_connections:
            username = self.player_usernames.get(player_id, "unknown")
            
            if websocket in self.active_connections[player_id]:
                self.active_connections[player_id].remove(websocket)
            
            # If no more connections for this player, cleanup
            if not self.active_connections[player_id]:
                del self.active_connections[player_id]
                if player_id in self.player_usernames:
                    del self.player_usernames[player_id]
                
                # Unsubscribe from player's personal channel
                if self.redis_manager:
                    await self.redis_manager.unsubscribe_from_channel(f"player:{player_id}")
                    await self.redis_manager.mark_player_disconnected(player_id)
                
                logger.info(f"All WebSockets disconnected: {username} (player_id={player_id})")
            else:
                logger.info(f"WebSocket disconnected: {username} (player_id={player_id}). Remaining connections: {len(self.active_connections[player_id])}")
    
    async def send_personal_message(self, player_id: int, message: dict):
        """Send a message to a specific player via Redis pub/sub."""
        if self.redis_manager:
            # Send locally first if player is connected to this worker
            if player_id in self.active_connections:
                await self._send_direct(player_id, message)
            else:
                # Publish to Redis (player might be on another worker)
                await self.redis_manager.publish_to_player(player_id, message)
        else:
            # Fallback to direct send (single worker mode)
            await self._send_direct(player_id, message)
    
    async def _send_direct(self, player_id: int, message: dict):
        """Directly send to local WebSocket connections."""
        if player_id in self.active_connections:
            connections = self.active_connections[player_id]
            disconnected_sockets = []
            
            # Inject unique message ID for tracing
            if "id" not in message:
                message["id"] = str(uuid.uuid4())
            
            for websocket in connections:
                try:
                    logger.debug(f"Using WS: Sending msg {message['id']} type={message.get('type')} to player {player_id}")
                    await websocket.send_json(message)
                except Exception as e:
                    logger.error(f"Failed to send message to player {player_id}: {e}")
                    disconnected_sockets.append(websocket)
            
            # Cleanup failed sockets
            for ws in disconnected_sockets:
                await self.disconnect(player_id, ws)
    
    async def broadcast(self, message: dict, exclude_player_id: Optional[int] = None):
        """Broadcast a message to all connected players via Redis."""
        if self.redis_manager:
            await self.redis_manager.publish_global_broadcast(message)
            
            # ALSO send to LOCAL connections immediately
            for player_id in list(self.active_connections.keys()):
                if player_id != exclude_player_id:
                    await self._send_direct(player_id, message)
        else:
            # Fallback: direct broadcast to local connections
            for player_id in list(self.active_connections.keys()):
                if player_id != exclude_player_id:
                    await self._send_direct(player_id, message)
    
    async def send_to_location(self, location_id: str, message: dict, exclude_player_id: Optional[int] = None):
        """Send a message to all players in a specific location via Redis pub/sub."""
        if self.redis_manager:
            # Use Redis pub/sub for cross-worker broadcast
            message_with_exclude = {
                **message,
                "exclude_player_id": exclude_player_id
            }
            await self.redis_manager.publish_to_location(location_id, message_with_exclude)
            
            # ALSO send to LOCAL connections immediately (don't wait for Redis roundtrip)
            player_ids = await self.redis_manager.get_players_in_location(location_id)
            for player_id in player_ids:
                if player_id == exclude_player_id:
                    continue
                if player_id in self.active_connections:
                    await self._send_direct(player_id, message)
        else:
            # Fallback: Query DB and send directly (single worker mode)
            from .. import database as db
            players_in_location = await db.get_players_in_location(location_id)
            
            active_players = [p for p in players_in_location if p['id'] in self.active_connections and p['id'] != exclude_player_id]
            if not active_players:
                return
            
            logger.info(f"Broadcasting to location {location_id}: {message.get('type')} (excluding player {exclude_player_id})")
            
            sent_count = 0
            for player in active_players:
                player_id = player['id']
                await self._send_direct(player_id, message)
                sent_count += 1
            
            logger.info(f"Sent {message.get('type')} to {sent_count} players")
    
    async def handle_redis_message(self, channel: str, data: dict):
        """
        Handle incoming Redis pub/sub messages and route to local WebSocket connections.
        This method is called by RedisManager when a message arrives on a subscribed channel.
        """
        try:
            # Extract message type and data
            message = {
                "type": data.get("type"),
                "data": data.get("data")
            }
            
            # Determine routing based on channel type
            if channel.startswith("player:"):
                # Personal message to specific player
                player_id = int(channel.split(":")[1])
                if player_id in self.active_connections:
                    await self._send_direct(player_id, message)
            
            elif channel.startswith("location:"):
                # Broadcast to all players in location (only local connections)
                location_id = channel.split(":")[1]
                exclude_player_id = data.get("exclude_player_id")
                
                # Get players from Redis location registry
                if self.redis_manager:
                    player_ids = await self.redis_manager.get_players_in_location(location_id)
                    
                    for player_id in player_ids:
                        if player_id == exclude_player_id:
                            continue
                        
                        # Only send if this worker has the connection
                        if player_id in self.active_connections:
                            await self._send_direct(player_id, message)
            
            elif channel == "game:broadcast":
                # Global broadcast to all local connections
                exclude_player_id = data.get("exclude_player_id")
                
                for player_id in list(self.active_connections.keys()):
                    if player_id != exclude_player_id:
                        await self._send_direct(player_id, message)
        
        except Exception as e:
            logger.error(f"Error handling Redis message on channel {channel}: {e}")
    
    def has_players_in_location(self, location_id: str) -> bool:
        """Check if there are any players with active connections in a specific location."""
        return len(self.active_connections) > 0
    
    def get_connected_count(self) -> int:
        """Get the number of active WebSocket connections."""
        return len(self.active_connections)


# Global connection manager instance
manager = ConnectionManager()
