city/server/main.py

import asyncio
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pathlib import Path
import time

from server.logger import logger
from server.websocket_manager import WebSocketManager
from server.game_state import GameState
from server.economy import EconomyEngine
from server.database import Database

# Global instances - can be overridden for testing
game_state = None
ws_manager = None
economy_engine = None
database = None

def initialize_components(test_db_path=None):
    """Initialize server components with optional test database"""
    global game_state, ws_manager, economy_engine, database
    
    # Initialize database first
    database = Database(test_db_path) if test_db_path else Database()
    database.init_db()
    
    # Create game state and load data from database
    game_state = GameState()
    game_state.load_state(database.load_game_state())
    
    # Now create WebSocketManager (it will rebuild nickname mapping from loaded players)
    ws_manager = WebSocketManager(game_state)
    economy_engine = EconomyEngine(game_state)
    
    return game_state, ws_manager, economy_engine, database

# --- HYBRID MODEL RE-IMPLEMENTED ---
last_economy_tick_time = time.time()
TICK_INTERVAL = 10  # seconds

async def trigger_economy_update_and_save():
    """Triggers an economy tick, broadcasts updates concurrently, saves, and resets the timer."""
    global last_economy_tick_time
    
    logger.debug("Triggering full economy update and save cycle...")
    start_time = time.perf_counter()

    economy_engine.tick()
    
    update_tasks = []
    for player_id, player in game_state.players.items():
        if player_id in ws_manager.active_connections:
            task = ws_manager.send_to_player(player_id, {
                "type": "player_stats_update",
                "player": player.to_dict()
            })
            update_tasks.append(task)
    
    if update_tasks:
        await asyncio.gather(*update_tasks)
    
    database.save_game_state(game_state)
    
    # Reset the global tick timer after any update
    last_economy_tick_time = time.time()
    
    duration = time.perf_counter() - start_time
    logger.info(f"Full economy update cycle completed in {duration:.4f} seconds.")

async def economy_loop():
    """Runs periodically to check if a passive economy update is needed."""
    global last_economy_tick_time
    while True:
        # Check if 10 seconds have passed since the last tick (from any source)
        if time.time() - last_economy_tick_time > TICK_INTERVAL:
            logger.info("Triggering timed economy update for idle players.")
            await trigger_economy_update_and_save()
        
        # Check frequently for responsiveness
        await asyncio.sleep(1)

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup and shutdown events."""
    global game_state, ws_manager, economy_engine, database
    
    # Initialize components if not already done (for testing)
    if not all([game_state, ws_manager, economy_engine, database]):
        initialize_components()
    
    logger.info("Server starting up...")
    # Database and game state are already initialized in initialize_components()
    
    # Start the hybrid economy loop
    task = asyncio.create_task(economy_loop())
    logger.info("Hybrid economy loop started.")
    
    yield
    
    logger.info("Server shutting down...")
    task.cancel()
    database.save_game_state(game_state)
    logger.info("Final game state saved.")

app = FastAPI(lifespan=lifespan)

app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
async def root():
    return FileResponse("static/index.html")

@app.websocket("/ws/{nickname}")
async def websocket_endpoint(websocket: WebSocket, nickname: str):
    await ws_manager.connect(websocket, nickname)
    try:
        player_id = await ws_manager.get_player_id(websocket)
        if player_id:
            player = game_state.get_or_create_player(nickname, player_id)
            
            # Trigger economy update on login to ensure player stats are current
            # This is especially important after server restarts
            logger.info(f"Player {player_id} connected, triggering economy update to sync stats.")
            await trigger_economy_update_and_save()
            
            # Send updated player data after economy sync
            updated_player = game_state.players.get(player_id, player)
            await websocket.send_json({
                "type": "init",
                "player": updated_player.to_dict(),
                "game_state": game_state.get_state()
            })
            while True:
                data = await websocket.receive_json()
                await handle_message(websocket, data)
    except WebSocketDisconnect:
        await ws_manager.disconnect(websocket)
    except Exception as e:
        player_id = await ws_manager.get_player_id(websocket)
        logger.error(f"An error occurred in the websocket endpoint for player {player_id}: {e}", exc_info=True)


async def handle_message(websocket: WebSocket, data: dict):
    """Handle incoming WebSocket messages."""
    msg_type = data.get("type")
    player_id = await ws_manager.get_player_id(websocket)

    if not player_id:
        logger.warning("Received message from an unknown player, ignoring.")
        return

    logger.debug(f"Received '{msg_type}' from player {player_id}")
    
    if msg_type == "cursor_move":
        await ws_manager.broadcast({"type": "cursor_move", "player_id": player_id, "x": data["x"], "y": data["y"]}, exclude=websocket)
    
    elif msg_type == "place_building":
        result = game_state.place_building(player_id, data["building_type"], data["x"], data["y"])
        if result["success"]:
            await ws_manager.broadcast({"type": "building_placed", "building": result["building"]})
            # --- CHANGE: Trigger an INSTANT economy update on financial action ---
            logger.info(f"Player {player_id} action triggered immediate economy update.")
            await trigger_economy_update_and_save()
        else:
            await websocket.send_json({"type": "error", "message": result["error"]})
    
    elif msg_type == "remove_building":
        result = game_state.remove_building(player_id, data["x"], data["y"])
        if result["success"]:
            await ws_manager.broadcast({"type": "building_removed", "x": data["x"], "y": data["y"]})
            # --- CHANGE: Trigger an INSTANT economy update on financial action ---
            logger.info(f"Player {player_id} action triggered immediate economy update.")
            await trigger_economy_update_and_save()
        else:
            await websocket.send_json({"type": "error", "message": result["error"]})

    elif msg_type == "edit_building":
        result = game_state.edit_building_name(player_id, data["x"], data["y"], data["name"])
        if result["success"]:
            await ws_manager.broadcast({"type": "building_updated", "x": data["x"], "y": data["y"], "name": data["name"]})
            database.save_game_state(game_state)
        else:
            await websocket.send_json({"type": "error", "message": result["error"]})
    
    elif msg_type == "chat":
        nickname = await ws_manager.get_nickname(websocket)
        await ws_manager.broadcast({"type": "chat", "nickname": nickname, "message": data["message"], "timestamp": data["timestamp"]})