city/server/database.py

import sqlite3
import json
from pathlib import Path
from typing import Optional, Dict, List
from server.logger import logger
import time

class Database:
    """Handles all database operations"""
    
    def __init__(self, db_path: str = "data/game.db"):
        self.db_path = db_path
        Path("data").mkdir(exist_ok=True)
    
    def _get_connection(self):
        """Get database connection"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn
    
    def init_db(self):
        """Initialize database tables"""
        logger.info(f"Initializing database at {self.db_path}...")
        with self._get_connection() as conn:
            conn.execute('''
                CREATE TABLE IF NOT EXISTS players (
                    player_id TEXT PRIMARY KEY,
                    nickname TEXT UNIQUE NOT NULL,
                    money INTEGER NOT NULL,
                    population INTEGER NOT NULL,
                    color TEXT NOT NULL,
                    last_online REAL NOT NULL
                )
            ''')
            conn.execute('''
                CREATE TABLE IF NOT EXISTS buildings (
                    x INTEGER NOT NULL,
                    y INTEGER NOT NULL,
                    type TEXT NOT NULL,
                    owner_id TEXT NOT NULL,
                    name TEXT,
                    placed_at REAL NOT NULL,
                    PRIMARY KEY (x, y),
                    FOREIGN KEY (owner_id) REFERENCES players (player_id)
                )
            ''')
            conn.commit()
        logger.info("Database initialized successfully.")
    
    def save_game_state(self, game_state):
        """Save complete game state to database"""
        logger.debug(f"Saving game state to database... ({len(game_state.players)} players, {len(game_state.buildings)} buildings)")
        start_time = time.perf_counter()
        with self._get_connection() as conn:
            player_data = [
                (p.player_id, p.nickname, p.money, p.population, p.color, p.last_online)
                for p in game_state.players.values()
            ]
            if player_data:
                conn.executemany('''
                    INSERT OR REPLACE INTO players (player_id, nickname, money, population, color, last_online)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', player_data)

            cursor = conn.execute('SELECT x, y FROM buildings')
            db_coords = {(row['x'], row['y']) for row in cursor}
            gs_coords = set(game_state.buildings.keys())

            coords_to_delete = db_coords - gs_coords
            if coords_to_delete:
                conn.executemany('DELETE FROM buildings WHERE x = ? AND y = ?', list(coords_to_delete))

            building_data = [
                (b.x, b.y, b.building_type.value, b.owner_id, b.name, b.placed_at)
                for b in game_state.buildings.values()
            ]
            if building_data:
                conn.executemany('''
                    INSERT OR REPLACE INTO buildings (x, y, type, owner_id, name, placed_at)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', building_data)
            
            conn.commit()
        duration = time.perf_counter() - start_time
        logger.debug(f"Game state saved to database in {duration:.4f} seconds.")
    
    def load_game_state(self) -> dict:
        """Load complete game state from database"""
        logger.info("Loading game state from database...")
        start_time = time.perf_counter()
        try:
            with self._get_connection() as conn:
                players, buildings = [], []
                
                cursor = conn.execute('SELECT * FROM players')
                for row in cursor:
                    players.append(dict(row))
                
                cursor = conn.execute('SELECT * FROM buildings')
                for row in cursor:
                    buildings.append(dict(row))
                
                duration = time.perf_counter() - start_time
                logger.info(f"Loaded {len(players)} players and {len(buildings)} buildings in {duration:.4f} seconds.")
                return {"players": players, "buildings": buildings}
        except Exception as e:
            logger.error(f"Error loading game state: {e}", exc_info=True)
            return {"players": [], "buildings": []}