import time import os from typing import Set from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy import ( MetaData, Table, Column, Integer, String, Boolean, ForeignKey, Float, UniqueConstraint, ) DB_USER, DB_PASS, DB_NAME, DB_HOST, DB_PORT = os.getenv("POSTGRES_USER"), os.getenv("POSTGRES_PASSWORD"), os.getenv("POSTGRES_DB"), os.getenv("POSTGRES_HOST"), os.getenv("POSTGRES_PORT") DATABASE_URL = f"postgresql+psycopg://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}" engine = create_async_engine(DATABASE_URL) metadata = MetaData() # ... (players, inventory, dropped_items tables are unchanged) ... players = Table( "players", metadata, Column("telegram_id", Integer, primary_key=True), Column("id", Integer, unique=True, autoincrement=True), # Web users ID Column("username", String(50), unique=True, nullable=True), # Web users username Column("password_hash", String(255), nullable=True), # Web users password hash Column("name", String, default="Survivor"), Column("hp", Integer, default=100), Column("max_hp", Integer, default=100), Column("stamina", Integer, default=20), Column("max_stamina", Integer, default=20), Column("strength", Integer, default=5), Column("agility", Integer, default=5), Column("endurance", Integer, default=5), Column("intellect", Integer, default=5), Column("location_id", String, default="start_point"), Column("is_dead", Boolean, default=False), Column("level", Integer, default=1), Column("xp", Integer, default=0), Column("unspent_points", Integer, default=0) ) inventory = Table("inventory", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("player_id", Integer, ForeignKey("players.telegram_id", ondelete="CASCADE")), Column("item_id", String), Column("quantity", Integer, default=1), Column("is_equipped", Boolean, default=False)) dropped_items = Table("dropped_items", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("item_id", String), Column("quantity", Integer, default=1), Column("location_id", String), Column("drop_timestamp", Float)) # Combat-related tables active_combats = Table( "active_combats", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("player_id", Integer, ForeignKey("players.telegram_id", ondelete="CASCADE"), unique=True), Column("npc_id", String, nullable=False), Column("npc_hp", Integer, nullable=False), Column("npc_max_hp", Integer, nullable=False), Column("turn", String, nullable=False), # "player" or "npc" Column("turn_started_at", Float, nullable=False), Column("player_status_effects", String, default=""), # JSON string Column("npc_status_effects", String, default=""), # JSON string Column("location_id", String, nullable=False), Column("from_wandering_enemy", Boolean, default=False), # If True, respawn on flee/death ) player_corpses = Table( "player_corpses", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("player_name", String, nullable=False), Column("location_id", String, nullable=False), Column("items", String, nullable=False), # JSON string of items Column("death_timestamp", Float, nullable=False), ) npc_corpses = Table( "npc_corpses", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("npc_id", String, nullable=False), Column("location_id", String, nullable=False), Column("loot_remaining", String, nullable=False), # JSON string Column("death_timestamp", Float, nullable=False), ) interactable_cooldowns = Table( "interactable_cooldowns", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("interactable_instance_id", String, nullable=False, unique=True), # Renamed for clarity Column("expiry_timestamp", Float, nullable=False), ) # Table to cache Telegram file IDs for images image_cache = Table( "image_cache", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("image_path", String, nullable=False, unique=True), # Local file path Column("telegram_file_id", String, nullable=False), # Telegram's file_id for reuse Column("uploaded_at", Float, nullable=False), ) # Wandering enemies table - managed by spawn system wandering_enemies = Table( "wandering_enemies", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("npc_id", String, nullable=False), Column("location_id", String, nullable=False), Column("spawn_timestamp", Float, nullable=False), Column("despawn_timestamp", Float, nullable=False), # When this enemy should despawn ) # Persistent status effects table player_status_effects = Table( "player_status_effects", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("player_id", Integer, ForeignKey("players.telegram_id", ondelete="CASCADE"), nullable=False), Column("effect_name", String(50), nullable=False), Column("effect_icon", String(10), nullable=False), Column("damage_per_tick", Integer, nullable=False, default=0), Column("ticks_remaining", Integer, nullable=False), Column("applied_at", Float, nullable=False), ) async def create_tables(): async with engine.begin() as conn: await conn.run_sync(metadata.create_all) # ... (All other database functions are unchanged except the cooldown ones) ... async def get_player(telegram_id: int = None, player_id: int = None, username: str = None): """Get player by telegram_id, player_id (web users), or username.""" async with engine.connect() as conn: if telegram_id is not None: result = await conn.execute(players.select().where(players.c.telegram_id == telegram_id)) elif player_id is not None: result = await conn.execute(players.select().where(players.c.id == player_id)) elif username is not None: result = await conn.execute(players.select().where(players.c.username == username)) else: return None row = result.first() return row._asdict() if row else None async def create_player(telegram_id: int = None, name: str = "Survivor", username: str = None, password_hash: str = None): """Create a player (Telegram or web user).""" async with engine.connect() as conn: values = { "name": name, "telegram_id": telegram_id, "username": username, "password_hash": password_hash, } result = await conn.execute(players.insert().values(**values)) await conn.commit() # For telegram users, the primary key is telegram_id # For web users, we need to get the auto-generated id if telegram_id: # Add starting inventory for Telegram users await conn.execute(inventory.insert().values(player_id=telegram_id, item_id="tattered_rucksack", is_equipped=True)) await conn.commit() # Return the created player if telegram_id: return await get_player(telegram_id=telegram_id) elif username: return await get_player(username=username) async def update_player(telegram_id: int = None, player_id: int = None, updates: dict = None): """Update player by telegram_id (Telegram users) or player_id (web users).""" if updates is None: updates = {} async with engine.connect() as conn: if telegram_id is not None: await conn.execute(players.update().where(players.c.telegram_id == telegram_id).values(**updates)) elif player_id is not None: await conn.execute(players.update().where(players.c.id == player_id).values(**updates)) else: raise ValueError("Must provide either telegram_id or player_id") await conn.commit() async def get_inventory(player_id: int): async with engine.connect() as conn: result = await conn.execute(inventory.select().where(inventory.c.player_id == player_id)) return [row._asdict() for row in result.fetchall()] async def get_inventory_item(item_db_id: int): async with engine.connect() as conn: result = await conn.execute(inventory.select().where(inventory.c.id == item_db_id)) row = result.first() return row._asdict() if row else None async def add_item_to_inventory(player_id: int, item_id: str, quantity: int = 1): async with engine.connect() as conn: result = await conn.execute(inventory.select().where(inventory.c.player_id == player_id, inventory.c.item_id == item_id)) existing_item = result.first() if existing_item: stmt = inventory.update().where(inventory.c.id == existing_item.id).values(quantity=inventory.c.quantity + quantity) else: stmt = inventory.insert().values(player_id=player_id, item_id=item_id, quantity=quantity) await conn.execute(stmt) await conn.commit() async def add_equipped_item_to_inventory(player_id: int, item_id: str) -> int: """Add a single equipped item to inventory and return its ID.""" async with engine.connect() as conn: stmt = inventory.insert().values( player_id=player_id, item_id=item_id, quantity=1, is_equipped=True ) result = await conn.execute(stmt) await conn.commit() return result.inserted_primary_key[0] async def update_inventory_item(item_db_id: int, quantity: int = None, is_equipped: bool = None): """Update inventory item properties.""" async with engine.connect() as conn: updates = {} if quantity is not None: updates['quantity'] = quantity if is_equipped is not None: updates['is_equipped'] = is_equipped if updates: stmt = inventory.update().where(inventory.c.id == item_db_id).values(**updates) await conn.execute(stmt) await conn.commit() async def remove_item_from_inventory(item_db_id: int, quantity: int = 1): async with engine.connect() as conn: result = await conn.execute(inventory.select().where(inventory.c.id == item_db_id)) item_data = result.first() if not item_data: return if item_data.quantity > quantity: stmt = inventory.update().where(inventory.c.id == item_db_id).values(quantity=inventory.c.quantity - quantity) else: stmt = inventory.delete().where(inventory.c.id == item_db_id) await conn.execute(stmt) await conn.commit() async def drop_item_to_world(item_id: str, quantity: int, location_id: str): """Drop item to world. Combines with existing stacks of same item in same location.""" async with engine.connect() as conn: # Check if this item already exists in this location result = await conn.execute( dropped_items.select().where( (dropped_items.c.item_id == item_id) & (dropped_items.c.location_id == location_id) ) ) existing_item = result.first() if existing_item: # Stack exists, add to it new_quantity = existing_item.quantity + quantity stmt = dropped_items.update().where(dropped_items.c.id == existing_item.id).values( quantity=new_quantity, drop_timestamp=time.time() # Update timestamp ) else: # Create new stack stmt = dropped_items.insert().values( item_id=item_id, quantity=quantity, location_id=location_id, drop_timestamp=time.time() ) await conn.execute(stmt) await conn.commit() async def get_dropped_items_in_location(location_id: str): async with engine.connect() as conn: result = await conn.execute(dropped_items.select().where(dropped_items.c.location_id == location_id).limit(10)) return [row._asdict() for row in result.fetchall()] async def get_dropped_item(dropped_item_id: int): async with engine.connect() as conn: result = await conn.execute(dropped_items.select().where(dropped_items.c.id == dropped_item_id)) row = result.first() return row._asdict() if row else None async def remove_dropped_item(dropped_item_id: int): async with engine.connect() as conn: await conn.execute(dropped_items.delete().where(dropped_items.c.id == dropped_item_id)) await conn.commit() async def update_dropped_item(dropped_item_id: int, new_quantity: int): """Update the quantity of a dropped item.""" async with engine.connect() as conn: stmt = dropped_items.update().where(dropped_items.c.id == dropped_item_id).values(quantity=new_quantity) await conn.execute(stmt) await conn.commit() async def remove_expired_dropped_items(timestamp_limit: float) -> int: async with engine.connect() as conn: stmt = dropped_items.delete().where(dropped_items.c.drop_timestamp < timestamp_limit) result = await conn.execute(stmt) await conn.commit() return result.rowcount async def regenerate_all_players_stamina() -> int: """ Regenerate stamina for all active players using a single optimized query. Recovery formula: - Base recovery: 1 stamina per cycle (5 minutes) - Endurance bonus: +1 stamina per 10 endurance points - Example: 5 endurance = 1 stamina, 15 endurance = 2 stamina, 25 endurance = 3 stamina - Only regenerates up to max_stamina - Only regenerates for living players PERFORMANCE: Single SQL query, scales to 100K+ players efficiently. """ from sqlalchemy import text async with engine.connect() as conn: # Single UPDATE query with database-side calculation # Much more efficient than fetching all players and updating individually stmt = text(""" UPDATE players SET stamina = LEAST( stamina + 1 + (endurance / 10), max_stamina ) WHERE is_dead = FALSE AND stamina < max_stamina """) result = await conn.execute(stmt) await conn.commit() return result.rowcount COOLDOWN_DURATION = 300 async def set_cooldown(instance_id: str): expiry_time = time.time() + COOLDOWN_DURATION async with engine.connect() as conn: update_stmt = interactable_cooldowns.update().where(interactable_cooldowns.c.interactable_instance_id == instance_id).values(expiry_timestamp=expiry_time) result = await conn.execute(update_stmt) if result.rowcount == 0: insert_stmt = interactable_cooldowns.insert().values(interactable_instance_id=instance_id, expiry_timestamp=expiry_time) await conn.execute(insert_stmt) await conn.commit() # --- Combat Functions --- async def create_combat(player_id: int, npc_id: str, npc_hp: int, npc_max_hp: int, location_id: str, from_wandering_enemy: bool = False): """Start a new combat encounter.""" async with engine.connect() as conn: stmt = active_combats.insert().values( player_id=player_id, npc_id=npc_id, npc_hp=npc_hp, npc_max_hp=npc_max_hp, turn="player", turn_started_at=time.time(), location_id=location_id, player_status_effects="[]", npc_status_effects="[]", from_wandering_enemy=from_wandering_enemy ) result = await conn.execute(stmt) await conn.commit() return result.inserted_primary_key[0] async def get_combat(player_id: int): """Get active combat for a player.""" async with engine.connect() as conn: stmt = active_combats.select().where(active_combats.c.player_id == player_id) result = await conn.execute(stmt) row = result.first() return row._asdict() if row else None async def update_combat(player_id: int, updates: dict): """Update combat state.""" async with engine.connect() as conn: stmt = active_combats.update().where(active_combats.c.player_id == player_id).values(**updates) await conn.execute(stmt) await conn.commit() async def end_combat(player_id: int): """Remove active combat.""" async with engine.connect() as conn: stmt = active_combats.delete().where(active_combats.c.player_id == player_id) await conn.execute(stmt) await conn.commit() async def get_all_idle_combats(idle_threshold: float): """Get all combats where the turn has been idle too long.""" async with engine.connect() as conn: stmt = active_combats.select().where(active_combats.c.turn_started_at < idle_threshold) result = await conn.execute(stmt) return [row._asdict() for row in result.fetchall()] async def create_player_corpse(player_name: str, location_id: str, items: str): """Create a player corpse bag.""" async with engine.connect() as conn: stmt = player_corpses.insert().values( player_name=player_name, location_id=location_id, items=items, death_timestamp=time.time() ) await conn.execute(stmt) await conn.commit() async def get_player_corpses_in_location(location_id: str): """Get all player corpses in a location.""" async with engine.connect() as conn: stmt = player_corpses.select().where(player_corpses.c.location_id == location_id) result = await conn.execute(stmt) return [row._asdict() for row in result.fetchall()] async def get_player_corpse(corpse_id: int): """Get a specific player corpse.""" async with engine.connect() as conn: stmt = player_corpses.select().where(player_corpses.c.id == corpse_id) result = await conn.execute(stmt) row = result.first() return row._asdict() if row else None async def update_player_corpse(corpse_id: int, items: str): """Update items in a player corpse.""" async with engine.connect() as conn: stmt = player_corpses.update().where(player_corpses.c.id == corpse_id).values(items=items) await conn.execute(stmt) await conn.commit() async def remove_player_corpse(corpse_id: int): """Remove a player corpse.""" async with engine.connect() as conn: stmt = player_corpses.delete().where(player_corpses.c.id == corpse_id) await conn.execute(stmt) await conn.commit() async def remove_expired_player_corpses(timestamp_limit: float) -> int: """Remove old player corpses.""" async with engine.connect() as conn: stmt = player_corpses.delete().where(player_corpses.c.death_timestamp < timestamp_limit) result = await conn.execute(stmt) await conn.commit() return result.rowcount async def create_npc_corpse(npc_id: str, location_id: str, loot_remaining: str): """Create an NPC corpse for scavenging.""" async with engine.connect() as conn: stmt = npc_corpses.insert().values( npc_id=npc_id, location_id=location_id, loot_remaining=loot_remaining, death_timestamp=time.time() ) result = await conn.execute(stmt) await conn.commit() return result.inserted_primary_key[0] async def get_npc_corpses_in_location(location_id: str): """Get all NPC corpses in a location.""" async with engine.connect() as conn: stmt = npc_corpses.select().where(npc_corpses.c.location_id == location_id) result = await conn.execute(stmt) return [row._asdict() for row in result.fetchall()] async def get_npc_corpse(corpse_id: int): """Get a specific NPC corpse.""" async with engine.connect() as conn: stmt = npc_corpses.select().where(npc_corpses.c.id == corpse_id) result = await conn.execute(stmt) row = result.first() return row._asdict() if row else None async def update_npc_corpse(corpse_id: int, loot_remaining: str): """Update loot in an NPC corpse.""" async with engine.connect() as conn: stmt = npc_corpses.update().where(npc_corpses.c.id == corpse_id).values(loot_remaining=loot_remaining) await conn.execute(stmt) await conn.commit() async def remove_npc_corpse(corpse_id: int): """Remove an NPC corpse.""" async with engine.connect() as conn: stmt = npc_corpses.delete().where(npc_corpses.c.id == corpse_id) await conn.execute(stmt) await conn.commit() async def remove_expired_npc_corpses(timestamp_limit: float) -> int: """Remove old NPC corpses.""" async with engine.connect() as conn: stmt = npc_corpses.delete().where(npc_corpses.c.death_timestamp < timestamp_limit) result = await conn.execute(stmt) await conn.commit() return result.rowcount async def get_cooldown(instance_id: str) -> int: async with engine.connect() as conn: stmt = interactable_cooldowns.select().where(interactable_cooldowns.c.interactable_instance_id == instance_id) result = await conn.execute(stmt) cooldown = result.first() if cooldown and cooldown.expiry_timestamp > time.time(): return int(cooldown.expiry_timestamp - time.time()) return 0 async def get_cooldowns_for_location(location_id: str) -> Set[str]: """Get all active cooldown instance IDs for a location by checking the prefix.""" async with engine.connect() as conn: stmt = interactable_cooldowns.select().where( interactable_cooldowns.c.interactable_instance_id.startswith(location_id + "_"), interactable_cooldowns.c.expiry_timestamp > time.time() ) result = await conn.execute(stmt) return {row.interactable_instance_id for row in result.fetchall()} # --- Image Cache Functions --- async def get_cached_image(image_path: str): """Get the Telegram file_id for a cached image.""" async with engine.connect() as conn: stmt = image_cache.select().where(image_cache.c.image_path == image_path) result = await conn.execute(stmt) row = result.first() return row.telegram_file_id if row else None async def cache_image(image_path: str, telegram_file_id: str): """Store a Telegram file_id for an image path.""" async with engine.connect() as conn: # Check if already exists stmt = image_cache.select().where(image_cache.c.image_path == image_path) result = await conn.execute(stmt) existing = result.first() if existing: # Update existing entry update_stmt = image_cache.update().where( image_cache.c.image_path == image_path ).values(telegram_file_id=telegram_file_id, uploaded_at=time.time()) await conn.execute(update_stmt) else: # Insert new entry insert_stmt = image_cache.insert().values( image_path=image_path, telegram_file_id=telegram_file_id, uploaded_at=time.time() ) await conn.execute(insert_stmt) await conn.commit() # --- Wandering Enemies Functions --- async def spawn_wandering_enemy(npc_id: str, location_id: str, lifetime_seconds: int = 600): """Spawn a wandering enemy at a location. Lifetime defaults to 10 minutes.""" async with engine.connect() as conn: current_time = time.time() despawn_time = current_time + lifetime_seconds await conn.execute(wandering_enemies.insert().values( npc_id=npc_id, location_id=location_id, spawn_timestamp=current_time, despawn_timestamp=despawn_time )) await conn.commit() async def get_wandering_enemies_in_location(location_id: str): """Get all active wandering enemies at a location.""" async with engine.connect() as conn: current_time = time.time() stmt = wandering_enemies.select().where( wandering_enemies.c.location_id == location_id, wandering_enemies.c.despawn_timestamp > current_time ) result = await conn.execute(stmt) return [row._asdict() for row in result.fetchall()] async def remove_wandering_enemy(enemy_id: int): """Remove a wandering enemy (when engaged in combat or manually despawned).""" async with engine.connect() as conn: await conn.execute(wandering_enemies.delete().where(wandering_enemies.c.id == enemy_id)) await conn.commit() async def cleanup_expired_wandering_enemies(): """Remove all expired wandering enemies.""" async with engine.connect() as conn: current_time = time.time() result = await conn.execute( wandering_enemies.delete().where(wandering_enemies.c.despawn_timestamp <= current_time) ) await conn.commit() return result.rowcount # Number of enemies despawned async def get_wandering_enemy_count_in_location(location_id: str) -> int: """Count active wandering enemies at a location.""" async with engine.connect() as conn: current_time = time.time() from sqlalchemy import func stmt = wandering_enemies.select().where( wandering_enemies.c.location_id == location_id, wandering_enemies.c.despawn_timestamp > current_time ) result = await conn.execute(stmt) return len(result.fetchall()) async def get_all_active_wandering_enemies(): """Get all active wandering enemies across all locations.""" async with engine.connect() as conn: current_time = time.time() stmt = wandering_enemies.select().where( wandering_enemies.c.despawn_timestamp > current_time ) result = await conn.execute(stmt) return [row._asdict() for row in result.fetchall()] # ============================================================================ # STATUS EFFECTS # ============================================================================ async def get_player_status_effects(player_id: int): """Get all active status effects for a player.""" async with engine.connect() as conn: stmt = player_status_effects.select().where( player_status_effects.c.player_id == player_id, player_status_effects.c.ticks_remaining > 0 ) result = await conn.execute(stmt) return [row._asdict() for row in result.fetchall()] async def add_status_effect(player_id: int, effect_name: str, effect_icon: str, damage_per_tick: int, ticks_remaining: int): """Add a new status effect to a player.""" async with engine.connect() as conn: await conn.execute( player_status_effects.insert().values( player_id=player_id, effect_name=effect_name, effect_icon=effect_icon, damage_per_tick=damage_per_tick, ticks_remaining=ticks_remaining, applied_at=time.time() ) ) await conn.commit() async def update_status_effect_ticks(effect_id: int, ticks_remaining: int): """Update the remaining ticks for a status effect.""" async with engine.connect() as conn: await conn.execute( player_status_effects.update().where( player_status_effects.c.id == effect_id ).values(ticks_remaining=ticks_remaining) ) await conn.commit() async def remove_status_effect(effect_id: int): """Remove a specific status effect.""" async with engine.connect() as conn: await conn.execute( player_status_effects.delete().where(player_status_effects.c.id == effect_id) ) await conn.commit() async def remove_all_status_effects(player_id: int): """Remove all status effects from a player.""" async with engine.connect() as conn: await conn.execute( player_status_effects.delete().where(player_status_effects.c.player_id == player_id) ) await conn.commit() async def remove_status_effects_by_name(player_id: int, effect_name: str, count: int = 1): """ Remove a specific number of status effects by name for a player. Used for treatment items that cure specific effects. Returns the number of effects actually removed. """ async with engine.connect() as conn: # Get the effects to remove stmt = player_status_effects.select().where( player_status_effects.c.player_id == player_id, player_status_effects.c.effect_name == effect_name, player_status_effects.c.ticks_remaining > 0 ).limit(count) result = await conn.execute(stmt) effects_to_remove = result.fetchall() # Remove them effect_ids = [row.id for row in effects_to_remove] if effect_ids: await conn.execute( player_status_effects.delete().where( player_status_effects.c.id.in_(effect_ids) ) ) await conn.commit() return len(effect_ids) async def get_all_players_with_status_effects(): """Get all player IDs that have active status effects (for background processing).""" async with engine.connect() as conn: from sqlalchemy import distinct stmt = player_status_effects.select().with_only_columns( distinct(player_status_effects.c.player_id) ).where(player_status_effects.c.ticks_remaining > 0) result = await conn.execute(stmt) return [row[0] for row in result.fetchall()] async def decrement_all_status_effect_ticks(): """ Decrement ticks for all active status effects and return affected player IDs. Used by background processor. """ async with engine.connect() as conn: # Get player IDs with effects before updating from sqlalchemy import distinct stmt = player_status_effects.select().with_only_columns( distinct(player_status_effects.c.player_id) ).where(player_status_effects.c.ticks_remaining > 0) result = await conn.execute(stmt) affected_players = [row[0] for row in result.fetchall()] # Decrement ticks await conn.execute( player_status_effects.update().where( player_status_effects.c.ticks_remaining > 0 ).values(ticks_remaining=player_status_effects.c.ticks_remaining - 1) ) # Remove expired effects await conn.execute( player_status_effects.delete().where(player_status_effects.c.ticks_remaining <= 0) ) await conn.commit() return affected_players