"""Inject realistic test runs + encounters into the dev database. Usage: python -m app.seeds.inject_test_data """ import asyncio import random from datetime import datetime, timedelta, timezone from sqlalchemy import delete, select from sqlalchemy.ext.asyncio import AsyncSession from app.core.database import async_session from app.models.encounter import Encounter from app.models.evolution import Evolution from app.models.game import Game from app.models.nuzlocke_run import NuzlockeRun from app.models.pokemon import Pokemon from app.models.route import Route random.seed(42) # reproducible data # --- Nicknames pool --- NICKNAMES = [ "Blaze", "Thunder", "Shadow", "Luna", "Spike", "Rex", "Cinder", "Misty", "Rocky", "Breeze", "Fang", "Nova", "Scout", "Atlas", "Pepper", "Storm", "Bandit", "Echo", "Maple", "Titan", "Ziggy", "Bolt", "Rusty", "Pearl", "Ivy", "Ghost", "Sunny", "Dash", "Ember", "Frost", "Jade", "Onyx", "Willow", "Tank", "Pip", "Mochi", "Salem", "Patches", "Bean", "Rocket", ] DEATH_CAUSES = [ "Crit from Gym Leader's ace", "Poison damage after battle", "Rival battle ambush", "Self-destruct from wild Geodude", "Underleveled against Elite Four", "Double battle gone wrong", "Missed the switch, one-shot by Hyper Beam", "Wild Wobbuffet trapped and KO'd", "Explosion from wild Graveler", "Unexpected crit from trainer's Pokemon", ] # --- Run definitions --- # Each run specifies a game slug, name, status, rules overrides, and # how far through the route list it should progress (as a fraction). RUN_DEFS = [ # Failed runs { "game_slug": "firered", "name": "Kanto Heartbreak", "status": "failed", "progress": 0.45, "rules": {"hardcoreMode": True, "setModeOnly": True}, "started_days_ago": 30, "ended_days_ago": 20, }, { "game_slug": "platinum", "name": "Sinnoh Disaster", "status": "failed", "progress": 0.30, "rules": {"pinwheelClause": False}, "started_days_ago": 60, "ended_days_ago": 45, }, # Completed runs { "game_slug": "emerald", "name": "Hoenn Champion Run", "status": "completed", "progress": 0.90, "rules": {"levelCaps": True}, "started_days_ago": 90, "ended_days_ago": 70, }, { "game_slug": "heartgold", "name": "Johto Victory Lap", "status": "completed", "progress": 0.85, "rules": {}, "started_days_ago": 120, "ended_days_ago": 95, }, # Active runs { "game_slug": "black", "name": "Unova Adventure", "status": "active", "progress": 0.35, "rules": {}, "started_days_ago": 5, "ended_days_ago": None, }, { "game_slug": "crystal", "name": "Crystal Nuzlocke", "status": "active", "progress": 0.20, "rules": {"hardcoreMode": True, "levelCaps": True, "setModeOnly": True}, "started_days_ago": 2, "ended_days_ago": None, }, ] # Default rules (matches frontend DEFAULT_RULES) DEFAULT_RULES = { "firstEncounterOnly": True, "permadeath": True, "nicknameRequired": True, "duplicatesClause": True, "shinyClause": True, "pinwheelClause": True, "hardcoreMode": False, "levelCaps": False, "setModeOnly": False, } async def get_game_by_slug(session: AsyncSession, slug: str) -> Game | None: result = await session.execute(select(Game).where(Game.slug == slug)) return result.scalar_one_or_none() async def get_leaf_routes(session: AsyncSession, game_id: int) -> list[Route]: """Get routes that can have encounters (no children).""" # Get all routes for the game result = await session.execute( select(Route) .where(Route.game_id == game_id) .order_by(Route.order) ) all_routes = result.scalars().all() parent_ids = {r.parent_route_id for r in all_routes if r.parent_route_id is not None} leaf_routes = [r for r in all_routes if r.id not in parent_ids] return leaf_routes async def get_encounterables( session: AsyncSession, game_id: int ) -> list[int]: """Get pokemon IDs that appear in route encounters for this game.""" from app.models.route_encounter import RouteEncounter result = await session.execute( select(RouteEncounter.pokemon_id) .join(Route, Route.id == RouteEncounter.route_id) .where(Route.game_id == game_id) .distinct() ) return [row[0] for row in result] async def get_evolution_map(session: AsyncSession) -> dict[int, list[int]]: """Return {from_pokemon_id: [to_pokemon_id, ...]} for all evolutions.""" result = await session.execute(select(Evolution.from_pokemon_id, Evolution.to_pokemon_id)) evo_map: dict[int, list[int]] = {} for from_id, to_id in result: evo_map.setdefault(from_id, []).append(to_id) return evo_map def pick_routes_for_run( leaf_routes: list[Route], progress: float ) -> list[Route]: """Pick a subset of leaf routes respecting one-per-group. For routes with a parent, only one sibling per parent_route_id is chosen. """ count = max(1, int(len(leaf_routes) * progress)) # Separate standalone and grouped routes standalone = [r for r in leaf_routes if r.parent_route_id is None] by_parent: dict[int, list[Route]] = {} for r in leaf_routes: if r.parent_route_id is not None: by_parent.setdefault(r.parent_route_id, []).append(r) # Build a pool: all standalone + one random child per group pool: list[Route] = list(standalone) for children in by_parent.values(): pool.append(random.choice(children)) random.shuffle(pool) # Take routes in order up to the target count pool.sort(key=lambda r: r.order) return pool[:count] def generate_encounter( run_id: int, route: Route, pokemon_ids: list[int], evo_map: dict[int, list[int]], used_pokemon: set[int], encounter_index: int, ) -> Encounter: """Generate a single encounter with realistic varied states.""" # Pick a pokemon, preferring ones not yet used (duplicates clause flavor) available = [p for p in pokemon_ids if p not in used_pokemon] if not available: available = pokemon_ids pokemon_id = random.choice(available) used_pokemon.add(pokemon_id) # Decide status with weighted distribution roll = random.random() if roll < 0.55: status = "caught" elif roll < 0.80: status = "fainted" else: status = "missed" nickname = None catch_level = None faint_level = None death_cause = None current_pokemon_id = None if status == "caught": nickname = NICKNAMES[encounter_index % len(NICKNAMES)] catch_level = random.randint(5, 55) # ~30% of caught pokemon die if random.random() < 0.30: faint_level = catch_level + random.randint(1, 15) death_cause = random.choice(DEATH_CAUSES) # ~25% of caught pokemon that are alive evolve if faint_level is None and random.random() < 0.25: evos = evo_map.get(pokemon_id, []) if evos: current_pokemon_id = random.choice(evos) elif status == "fainted": catch_level = random.randint(5, 45) return Encounter( run_id=run_id, route_id=route.id, pokemon_id=pokemon_id, nickname=nickname, status=status, catch_level=catch_level, faint_level=faint_level, death_cause=death_cause, current_pokemon_id=current_pokemon_id, ) async def inject(): """Clear existing runs and inject test data.""" print("Injecting test data...") async with async_session() as session: async with session.begin(): # Clear existing runs and encounters await session.execute(delete(Encounter)) await session.execute(delete(NuzlockeRun)) print("Cleared existing runs and encounters") evo_map = await get_evolution_map(session) now = datetime.now(timezone.utc) total_runs = 0 total_encounters = 0 for run_def in RUN_DEFS: game = await get_game_by_slug(session, run_def["game_slug"]) if game is None: print(f" Warning: game '{run_def['game_slug']}' not found, skipping") continue # Build rules rules = {**DEFAULT_RULES, **run_def["rules"]} # Compute dates started_at = now - timedelta(days=run_def["started_days_ago"]) completed_at = None if run_def["ended_days_ago"] is not None: completed_at = now - timedelta(days=run_def["ended_days_ago"]) run = NuzlockeRun( game_id=game.id, name=run_def["name"], status=run_def["status"], rules=rules, started_at=started_at, completed_at=completed_at, ) session.add(run) await session.flush() # get run.id # Get routes and pokemon for this game leaf_routes = await get_leaf_routes(session, game.id) pokemon_ids = await get_encounterables(session, game.id) if not leaf_routes or not pokemon_ids: print(f" {run_def['name']}: no routes or pokemon, skipping encounters") total_runs += 1 continue chosen_routes = pick_routes_for_run(leaf_routes, run_def["progress"]) used_pokemon: set[int] = set() run_encounters = 0 for i, route in enumerate(chosen_routes): enc = generate_encounter( run.id, route, pokemon_ids, evo_map, used_pokemon, i ) session.add(enc) run_encounters += 1 total_runs += 1 total_encounters += run_encounters print( f" {run_def['name']} ({game.name}, {run_def['status']}): " f"{run_encounters} encounters across {len(chosen_routes)} routes" ) print(f"\nCreated {total_runs} runs with {total_encounters} total encounters") print("Test data injection complete!") if __name__ == "__main__": asyncio.run(inject())