city/server/game_state.py

from typing import Dict, List, Optional, Set, Tuple
from server.models import Player, Building, BuildingType, BUILDING_CONFIGS
import time

class GameState:
    """Manages the complete game state"""
    
    def __init__(self):
        self.players: Dict[str, Player] = {}
        self.buildings: Dict[Tuple[int, int], Building] = {}
        self.road_network: Set[Tuple[int, int]] = set()
        self.connected_zones: List[Set[Tuple[int, int]]] = []
    
    def get_or_create_player(self, nickname: str, player_id: str) -> Player:
        """Get existing player or create new one"""
        if player_id in self.players:
            player = self.players[player_id]
            player.is_online = True
            player.last_online = time.time()
            return player
        
        player = Player(
            player_id=player_id,
            nickname=nickname,
            last_online=time.time()
        )
        self.players[player_id] = player
        return player
    
    def place_building(self, player_id: str, building_type: str, x: int, y: int) -> dict:
        """Place a building on the map"""
        # Check if tile is occupied
        if (x, y) in self.buildings:
            return {"success": False, "error": "Tile already occupied"}
        
        # Get player
        player = self.players.get(player_id)
        if not player:
            return {"success": False, "error": "Player not found"}
        
        # Get building config
        try:
            b_type = BuildingType(building_type)
            config = BUILDING_CONFIGS[b_type]
        except (ValueError, KeyError):
            return {"success": False, "error": "Invalid building type"}
        
        # Check if player can afford
        if not player.can_afford(config.cost):
            return {"success": False, "error": "Not enough money"}
        
        # Check requirements
        if config.requires_population > player.population:
            return {"success": False, "error": f"Requires {config.requires_population} population"}
        
        if config.power_required and not self._has_power_plant(player_id):
            return {"success": False, "error": "Requires power plant"}
        
        # Place building
        building = Building(
            building_type=b_type,
            x=x,
            y=y,
            owner_id=player_id,
            placed_at=time.time()
        )
        
        self.buildings[(x, y)] = building
        player.deduct_money(config.cost)
        
        # Update road network if it's a road
        if b_type == BuildingType.ROAD:
            self.road_network.add((x, y))
            self._update_connected_zones()
        
        return {"success": True, "building": building.to_dict()}
    
    def remove_building(self, player_id: str, x: int, y: int) -> dict:
        """Remove a building"""
        building = self.buildings.get((x, y))
        
        if not building:
            return {"success": False, "error": "No building at this location"}
        
        if building.owner_id != player_id:
            return {"success": False, "error": "You don't own this building"}
        
        # Remove building
        del self.buildings[(x, y)]
        
        # Update road network if it was a road
        if building.building_type == BuildingType.ROAD:
            self.road_network.discard((x, y))
            self._update_connected_zones()
        
        return {"success": True}
    
    def edit_building_name(self, player_id: str, x: int, y: int, name: str) -> dict:
        """Edit building name"""
        building = self.buildings.get((x, y))
        
        if not building:
            return {"success": False, "error": "No building at this location"}
        
        if building.owner_id != player_id:
            return {"success": False, "error": "You don't own this building"}
        
        building.name = name
        return {"success": True}
    
    def _has_power_plant(self, player_id: str) -> bool:
        """Check if player has a power plant"""
        for building in self.buildings.values():
            if (building.owner_id == player_id and 
                building.building_type == BuildingType.POWER_PLANT):
                return True
        return False
    
    def _update_connected_zones(self):
        """Update connected zones based on road network using flood fill"""
        if not self.road_network:
            self.connected_zones = []
            return
        
        visited = set()
        self.connected_zones = []
        
        for road_pos in self.road_network:
            if road_pos in visited:
                continue
            
            # Flood fill to find connected zone
            zone = set()
            stack = [road_pos]
            
            while stack:
                pos = stack.pop()
                if pos in visited:
                    continue
                
                visited.add(pos)
                zone.add(pos)
                
                # Check adjacent positions
                x, y = pos
                for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
                    adj_pos = (x + dx, y + dy)
                    if adj_pos in self.road_network and adj_pos not in visited:
                        stack.append(adj_pos)
            
            self.connected_zones.append(zone)
    
    def get_building_zone_size(self, x: int, y: int) -> int:
        """Get the size of the connected zone for a building"""
        # Find adjacent roads
        for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (-1, 1), (1, -1), (-1, -1)]:
            road_pos = (x + dx, y + dy)
            if road_pos in self.road_network:
                # Find which zone this road belongs to
                for zone in self.connected_zones:
                    if road_pos in zone:
                        return len(zone)
        return 0
    
    def get_player_buildings(self, player_id: str) -> List[Building]:
        """Get all buildings owned by a player"""
        return [b for b in self.buildings.values() if b.owner_id == player_id]
    
    def get_state(self) -> dict:
        """Get complete game state for broadcasting"""
        return {
            "players": {pid: p.to_dict() for pid, p in self.players.items()},
            "buildings": {f"{x},{y}": b.to_dict() for (x, y), b in self.buildings.items()}
        }
    
    def load_state(self, state: dict):
        """Load game state from database"""
        if not state:
            return
        
        # Load players
        for player_data in state.get("players", []):
            player = Player(**player_data)
            player.is_online = False
            self.players[player.player_id] = player
        
        # Load buildings
        for building_data in state.get("buildings", []):
            building = Building(
                building_type=BuildingType(building_data["type"]),
                x=building_data["x"],
                y=building_data["y"],
                owner_id=building_data["owner_id"],
                name=building_data.get("name"),
                placed_at=building_data.get("placed_at", 0.0)
            )
            self.buildings[(building.x, building.y)] = building
            
            if building.building_type == BuildingType.ROAD:
                self.road_network.add((building.x, building.y))
        
        self._update_connected_zones()