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 collections.abc import Sequence
from pathlib import Path

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision: str = "d3e4f5a6b7c8"
down_revision: str | Sequence[str] | None = "c2d3e4f5a6b7"
branch_labels: str | Sequence[str] | None = None
depends_on: 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")