nuzlocke-tracker/backend/src/app/alembic/versions/d3e4f5a6b7c8_add_version_groups.py

"""add version groups

Revision ID: d3e4f5a6b7c8
Revises: c2d3e4f5a6b7
Create Date: 2026-02-08 14:00:00.000000

"""
import json
from pathlib import Path
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'd3e4f5a6b7c8'
down_revision: Union[str, Sequence[str], None] = 'c2d3e4f5a6b7'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # 1. Create version_groups table
    op.create_table(
        'version_groups',
        sa.Column('id', sa.Integer(), primary_key=True),
        sa.Column('name', sa.String(100), nullable=False),
        sa.Column('slug', sa.String(100), nullable=False, unique=True),
    )

    # 2. Populate version groups from seed data
    vg_json_path = Path(__file__).resolve().parents[2] / "seeds" / "version_groups.json"
    with open(vg_json_path) as f:
        vg_data = json.load(f)

    conn = op.get_bind()
    vg_table = sa.table(
        'version_groups',
        sa.column('id', sa.Integer),
        sa.column('name', sa.String),
        sa.column('slug', sa.String),
    )

    # Build slug -> id mapping and game_slug -> vg_id mapping
    slug_to_vg_id = {}
    game_slug_to_vg_id = {}
    for vg_idx, (vg_slug, vg_info) in enumerate(vg_data.items(), start=1):
        vg_id = vg_idx
        # Use the slug as a readable name (e.g., "red-blue" -> "Red / Blue")
        vg_name = " / ".join(
            g["name"].replace("Pokemon ", "")
            for g in vg_info["games"].values()
        )
        conn.execute(vg_table.insert().values(id=vg_id, name=vg_name, slug=vg_slug))
        slug_to_vg_id[vg_slug] = vg_id
        for game_slug in vg_info["games"]:
            game_slug_to_vg_id[game_slug] = vg_id

    # 3. Add version_group_id to games (nullable initially)
    op.add_column('games', sa.Column('version_group_id', sa.Integer(),
                                     sa.ForeignKey('version_groups.id'), nullable=True))
    op.create_index('ix_games_version_group_id', 'games', ['version_group_id'])

    # Populate games.version_group_id from the mapping
    games_table = sa.table(
        'games',
        sa.column('id', sa.Integer),
        sa.column('slug', sa.String),
        sa.column('version_group_id', sa.Integer),
    )
    rows = conn.execute(sa.select(games_table.c.id, games_table.c.slug)).fetchall()
    for game_id, game_slug in rows:
        vg_id = game_slug_to_vg_id.get(game_slug)
        if vg_id is not None:
            conn.execute(
                games_table.update()
                .where(games_table.c.id == game_id)
                .values(version_group_id=vg_id)
            )

    # 4. Add game_id to route_encounters (nullable initially), populate from routes.game_id
    op.add_column('route_encounters', sa.Column('game_id', sa.Integer(),
                                                sa.ForeignKey('games.id'), nullable=True))
    op.create_index('ix_route_encounters_game_id', 'route_encounters', ['game_id'])

    routes_table = sa.table(
        'routes',
        sa.column('id', sa.Integer),
        sa.column('name', sa.String),
        sa.column('game_id', sa.Integer),
    )
    re_table = sa.table(
        'route_encounters',
        sa.column('id', sa.Integer),
        sa.column('route_id', sa.Integer),
        sa.column('game_id', sa.Integer),
    )
    # Populate route_encounters.game_id from routes.game_id via join
    conn.execute(
        re_table.update()
        .where(re_table.c.route_id == routes_table.c.id)
        .values(game_id=routes_table.c.game_id)
    )

    # 5. Drop old unique constraint on route_encounters, add new one with game_id
    op.drop_constraint('uq_route_pokemon_method', 'route_encounters', type_='unique')
    op.create_unique_constraint(
        'uq_route_pokemon_method_game', 'route_encounters',
        ['route_id', 'pokemon_id', 'encounter_method', 'game_id']
    )

    # 6. Deduplicate routes within version groups
    # For multi-game version groups, keep routes from the lowest game_id (canonical)
    # and re-point route_encounters, encounters, and boss_battles to canonical routes

    encounters_table = sa.table(
        'encounters',
        sa.column('id', sa.Integer),
        sa.column('route_id', sa.Integer),
    )
    boss_battles_table = sa.table(
        'boss_battles',
        sa.column('id', sa.Integer),
        sa.column('game_id', sa.Integer),
        sa.column('after_route_id', sa.Integer),
    )

    # Get all version groups that have more than one game
    for vg_slug, vg_info in vg_data.items():
        game_slugs = list(vg_info["games"].keys())
        if len(game_slugs) <= 1:
            continue

        vg_id = slug_to_vg_id[vg_slug]

        # Get game IDs for this version group, ordered
        game_rows = conn.execute(
            sa.select(games_table.c.id, games_table.c.slug)
            .where(games_table.c.version_group_id == vg_id)
            .order_by(games_table.c.id)
        ).fetchall()

        if len(game_rows) <= 1:
            continue

        canonical_game_id = game_rows[0][0]
        non_canonical_game_ids = [r[0] for r in game_rows[1:]]

        # Get canonical routes (by name)
        canonical_routes = conn.execute(
            sa.select(routes_table.c.id, routes_table.c.name)
            .where(routes_table.c.game_id == canonical_game_id)
        ).fetchall()
        canonical_name_to_id = {name: rid for rid, name in canonical_routes}

        # For each non-canonical game, re-point references to canonical routes
        for nc_game_id in non_canonical_game_ids:
            nc_routes = conn.execute(
                sa.select(routes_table.c.id, routes_table.c.name)
                .where(routes_table.c.game_id == nc_game_id)
            ).fetchall()

            for old_route_id, route_name in nc_routes:
                canonical_id = canonical_name_to_id.get(route_name)
                if canonical_id is None:
                    continue

                # Re-point route_encounters
                conn.execute(
                    re_table.update()
                    .where(re_table.c.route_id == old_route_id)
                    .values(route_id=canonical_id)
                )

                # Re-point encounters
                conn.execute(
                    encounters_table.update()
                    .where(encounters_table.c.route_id == old_route_id)
                    .values(route_id=canonical_id)
                )

                # Re-point boss_battles.after_route_id
                conn.execute(
                    boss_battles_table.update()
                    .where(boss_battles_table.c.after_route_id == old_route_id)
                    .values(after_route_id=canonical_id)
                )

        # Delete non-canonical routes (children first due to parent FK)
        # First delete child routes
        conn.execute(
            sa.text(
                "DELETE FROM routes WHERE parent_route_id IS NOT NULL AND game_id IN :nc_ids"
            ).bindparams(sa.bindparam('nc_ids', expanding=True)),
            {"nc_ids": non_canonical_game_ids}
        )
        # Then delete parent routes
        conn.execute(
            sa.text(
                "DELETE FROM routes WHERE game_id IN :nc_ids"
            ).bindparams(sa.bindparam('nc_ids', expanding=True)),
            {"nc_ids": non_canonical_game_ids}
        )

    # 7. Add version_group_id to routes (nullable), populate from games.version_group_id
    op.add_column('routes', sa.Column('version_group_id', sa.Integer(),
                                      sa.ForeignKey('version_groups.id'), nullable=True))
    op.create_index('ix_routes_version_group_id', 'routes', ['version_group_id'])

    # Need to re-declare routes_table with version_group_id
    routes_table_v2 = sa.table(
        'routes',
        sa.column('id', sa.Integer),
        sa.column('name', sa.String),
        sa.column('game_id', sa.Integer),
        sa.column('version_group_id', sa.Integer),
    )

    # Populate routes.version_group_id from the game's version_group_id
    conn.execute(
        routes_table_v2.update()
        .where(routes_table_v2.c.game_id == games_table.c.id)
        .values(version_group_id=games_table.c.version_group_id)
    )

    # 8. Drop routes.game_id, drop old unique constraint, add new one
    op.drop_constraint('uq_routes_game_name', 'routes', type_='unique')
    op.drop_index('ix_routes_game_id', 'routes')
    op.drop_column('routes', 'game_id')
    op.create_unique_constraint(
        'uq_routes_version_group_name', 'routes',
        ['version_group_id', 'name']
    )

    # 9. Add version_group_id to boss_battles (nullable), populate from games.version_group_id
    op.add_column('boss_battles', sa.Column('version_group_id', sa.Integer(),
                                            sa.ForeignKey('version_groups.id'), nullable=True))
    op.create_index('ix_boss_battles_version_group_id', 'boss_battles', ['version_group_id'])

    bb_table_v2 = sa.table(
        'boss_battles',
        sa.column('id', sa.Integer),
        sa.column('game_id', sa.Integer),
        sa.column('version_group_id', sa.Integer),
    )

    conn.execute(
        bb_table_v2.update()
        .where(bb_table_v2.c.game_id == games_table.c.id)
        .values(version_group_id=games_table.c.version_group_id)
    )

    # 10. Drop boss_battles.game_id
    op.drop_index('ix_boss_battles_game_id', 'boss_battles')
    op.drop_column('boss_battles', 'game_id')

    # 11. Make columns non-nullable
    op.alter_column('route_encounters', 'game_id', nullable=False)
    op.alter_column('routes', 'version_group_id', nullable=False)
    op.alter_column('boss_battles', 'version_group_id', nullable=False)
    op.alter_column('games', 'version_group_id', nullable=False)


def downgrade() -> None:
    # This migration is not reversible in a meaningful way due to data deduplication
    raise NotImplementedError("Cannot downgrade: route deduplication is not reversible")