echoes-of-the-ash/bot/database.py

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("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
)

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):
    async with engine.connect() as conn:
        result = await conn.execute(players.select().where(players.c.telegram_id == telegram_id))
        row = result.first()
        return row._asdict() if row else None
async def create_player(telegram_id: int, name: str):
    async with engine.connect() as conn:
        await conn.execute(players.insert().values(telegram_id=telegram_id, name=name))
        await conn.execute(inventory.insert().values(player_id=telegram_id, item_id="tattered_rucksack", is_equipped=True))
        await conn.commit()
    return await get_player(telegram_id)
async def update_player(telegram_id: int, updates: dict):
    async with engine.connect() as conn:
        await conn.execute(players.update().where(players.c.telegram_id == telegram_id).values(**updates))
        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()]