"""add reference_sites and migrate references.site to FK

Revision ID: d7e3b19f0a2c
Revises: c3f8a1e924d5
Create Date: 2026-06-09
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import table, column, select

revision = "d7e3b19f0a2c"
down_revision = "c3f8a1e924d5"
branch_labels = None
depends_on = None


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

    # 2. Add site_id FK column (nullable)
    op.add_column(
        "references",
        sa.Column("site_id", sa.Integer(),
                  sa.ForeignKey("reference_sites.id", ondelete="SET NULL"),
                  nullable=True),
    )

    # 3. Migrate existing site string data
    conn = op.get_bind()
    refs_t  = table("references",      column("id", sa.Integer), column("site", sa.String), column("site_id", sa.Integer))
    sites_t = table("reference_sites", column("id", sa.Integer), column("name", sa.String))

    # Collect distinct non-null site strings
    existing = conn.execute(
        select(refs_t.c.site).where(refs_t.c.site.isnot(None)).distinct()
    ).fetchall()

    for (name,) in existing:
        # Insert into reference_sites (ignore duplicates)
        result = conn.execute(
            sites_t.insert().returning(sites_t.c.id).values(name=name)
        )
        site_id = result.fetchone()[0]
        # Update references rows with this site string
        conn.execute(
            refs_t.update().where(refs_t.c.site == name).values(site_id=site_id)
        )

    # 4. Drop old site string column
    op.drop_column("references", "site")


def downgrade() -> None:
    op.add_column("references", sa.Column("site", sa.String(), nullable=True))
    # best-effort: restore site name from FK
    conn = op.get_bind()
    refs_t  = table("references",      column("id", sa.Integer), column("site", sa.String), column("site_id", sa.Integer))
    sites_t = table("reference_sites", column("id", sa.Integer), column("name", sa.String))
    rows = conn.execute(select(refs_t.c.id, sites_t.c.name).join(sites_t, refs_t.c.site_id == sites_t.c.id)).fetchall()
    for ref_id, name in rows:
        conn.execute(refs_t.update().where(refs_t.c.id == ref_id).values(site=name))
    op.drop_column("references", "site_id")
    op.drop_table("reference_sites")