"""add account_id to content tables and backfill from _team_account_mapping

Revision ID: 019
Revises: 018
Create Date: 2026-02-07

Uses the _team_account_mapping table from migration 018 for deterministic
FK backfill instead of non-deterministic LIMIT 1 subqueries.
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID


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

# Tables that have team_id and need account_id
CONTENT_TABLES = ['trees', 'step_library', 'tree_categories', 'tree_tags', 'step_categories']


def upgrade() -> None:
    conn = op.get_bind()

    for table in CONTENT_TABLES:
        # Add account_id column
        op.add_column(table, sa.Column('account_id', UUID(as_uuid=True), nullable=True))
        op.create_index(f'ix_{table}_account_id', table, ['account_id'])

        # Backfill from mapping table (deterministic)
        conn.execute(sa.text(f"""
            UPDATE {table} SET account_id = m.account_id
            FROM _team_account_mapping m
            WHERE {table}.team_id = m.team_id
        """))

        # Validate: no rows with team_id but missing account_id
        orphaned = conn.execute(sa.text(f"""
            SELECT COUNT(*) FROM {table}
            WHERE team_id IS NOT NULL AND account_id IS NULL
        """)).scalar()
        if orphaned > 0:
            raise RuntimeError(
                f"Migration 019 failed: {table} has {orphaned} rows with team_id but no account_id"
            )


def downgrade() -> None:
    for table in reversed(CONTENT_TABLES):
        op.drop_index(f'ix_{table}_account_id', table_name=table)
        op.drop_column(table, 'account_id')