city/server/websocket_manager.py

from fastapi import WebSocket
from typing import Dict, List
import uuid
import time
from server.game_state import GameState
from server.logger import logger

class WebSocketManager:
    """Manages WebSocket connections for multiplayer"""
    
    def __init__(self, game_state: GameState):
        self.active_connections: Dict[str, List[WebSocket]] = {}
        self.player_nicknames: Dict[str, str] = {}
        self.nickname_to_id: Dict[str, str] = {}
        self.game_state = game_state
        
        # Rebuild nickname-to-ID mapping from game state
        self._rebuild_nickname_mapping()
    
    def _rebuild_nickname_mapping(self):
        """Rebuild the nickname-to-ID mapping from existing players in game state."""
        for player_id, player in self.game_state.players.items():
            self.nickname_to_id[player.nickname] = player_id
        
        if self.nickname_to_id:
            logger.info(f"Rebuilt nickname-to-ID mapping for {len(self.nickname_to_id)} players: {list(self.nickname_to_id.keys())}")
    
    async def connect(self, websocket: WebSocket, nickname: str):
        """Connect a new player, allowing multiple connections per nickname."""
        logger.debug(f"Connection attempt from nickname '{nickname}'.")
        
        await websocket.accept()
        
        if nickname in self.nickname_to_id:
            player_id = self.nickname_to_id[nickname]
        else:
            player_id = str(uuid.uuid4())
            self.nickname_to_id[nickname] = player_id
            logger.info(f"New player '{nickname}' assigned ID {player_id}.")
        
        player = self.game_state.get_or_create_player(nickname, player_id)
        
        if not self.active_connections.get(player_id):
            player.is_online = True
            self.active_connections[player_id] = []
            logger.info(f"Player '{nickname}' ({player_id}) is now online.")
            await self.broadcast({
                "type": "player_joined",
                "player_id": player_id,
                "nickname": nickname
            })

        self.active_connections[player_id].append(websocket)
        self.player_nicknames[player_id] = nickname
        logger.debug(f"Added new websocket for '{nickname}'. Total connections for user: {len(self.active_connections[player_id])}.")
    
    async def disconnect(self, websocket: WebSocket):
        """Disconnect a player's specific websocket instance."""
        player_id = await self.get_player_id(websocket)
        
        if player_id and player_id in self.active_connections:
            self.active_connections[player_id].remove(websocket)
            logger.debug(f"Removed websocket for player {player_id}. Remaining connections: {len(self.active_connections[player_id])}.")

            if not self.active_connections[player_id]:
                del self.active_connections[player_id]
                nickname = self.player_nicknames.pop(player_id, None)

                # --- REVERTED CHANGE: The following lines that set a player
                # --- to 'offline' have been removed to restore original economy.
                # player = self.game_state.players.get(player_id)
                # if player:
                #     player.is_online = False
                #     player.last_online = time.time()
                
                logger.info(f"Player '{nickname}' ({player_id}) last connection closed.")
                await self.broadcast({
                    "type": "player_left",
                    "player_id": player_id,
                    "nickname": nickname
                })
    
    async def send_to_player(self, player_id: str, message: dict):
        """Send a message to all active connections for a specific player."""
        connections = self.active_connections.get(player_id, [])
        if not connections: return
        
        logger.debug(f"Sending '{message['type']}' to player {player_id} ({len(connections)} connections).")
        disconnected_sockets = []

        for websocket in connections:
            try:
                await websocket.send_json(message)
            except Exception:
                disconnected_sockets.append(websocket)
        
        for ws in disconnected_sockets:
            logger.warning(f"Found dead socket for player {player_id} during send. Cleaning up.")
            if ws in self.active_connections.get(player_id, []):
                self.active_connections[player_id].remove(ws)

    async def broadcast(self, message: dict, exclude: WebSocket = None):
        """Broadcast message to all connected players and all their connections."""
        msg_type = message.get("type", "unknown")
        logger.debug(f"Broadcasting '{msg_type}' to all active connections.")
        all_sockets = []
        for ws_list in self.active_connections.values():
            all_sockets.extend(ws_list)

        for websocket in all_sockets:
            if websocket == exclude:
                continue
            
            try:
                await websocket.send_json(message)
            except Exception:
                pass
    
    async def get_player_id(self, websocket: WebSocket) -> str:
        """Get player ID from a specific websocket instance."""
        for player_id, ws_list in self.active_connections.items():
            if websocket in ws_list:
                return player_id
        return None
    
    async def get_nickname(self, websocket: WebSocket) -> str:
        """Get nickname from websocket."""
        player_id = await self.get_player_id(websocket)
        return self.player_nicknames.get(player_id, "Unknown")