"""set NOT NULL on all previously-nullable account_id columns

Revision ID: 174f442795b7
Revises: 3a40fe11b427
Create Date: 2026-04-09 00:00:00.000000

All tables in this migration had account_id set to nullable previously.
Task 9 (create_global_content_tables) cleared all NULL rows.
This migration enforces the NOT NULL constraint.
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa

revision: str = '174f442795b7'
down_revision: Union[str, None] = '3a40fe11b427'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # tree_embeddings: backfill from trees (must happen before SET NOT NULL)
    op.execute("""
        UPDATE tree_embeddings te
        SET account_id = t.account_id
        FROM trees t
        WHERE te.tree_id = t.id
          AND te.account_id IS NULL
    """)

    # feedback: backfill from users
    op.execute("""
        UPDATE feedback f
        SET account_id = u.account_id
        FROM users u
        WHERE f.user_id = u.id
          AND f.account_id IS NULL
    """)

    # Verify ALL tables before touching any SET NOT NULL
    tables_with_account_id = [
        'users', 'trees', 'tree_categories', 'tree_tags',
        'step_categories', 'step_library', 'tree_embeddings', 'feedback',
    ]
    for table in tables_with_account_id:
        result = op.get_bind().execute(
            sa.text(f"SELECT COUNT(*) FROM {table} WHERE account_id IS NULL")
        )
        count = result.scalar()
        if count > 0:
            raise RuntimeError(
                f"ROLLBACK: {count} NULL account_id rows in {table}. "
                "Run Task 9 (create_global_content_tables) first, or "
                "manually backfill/delete orphaned rows."
            )

    # SET NOT NULL on all
    for table in tables_with_account_id:
        op.alter_column(table, 'account_id', nullable=False)

    # Create indexes where they don't already exist
    new_indexes = [
        ('tree_embeddings', 'ix_tree_embeddings_account_id'),
        ('feedback', 'ix_feedback_account_id'),
    ]
    for table, index_name in new_indexes:
        result = op.get_bind().execute(sa.text(
            f"SELECT 1 FROM pg_indexes WHERE tablename='{table}' AND indexname='{index_name}'"
        ))
        if not result.fetchone():
            op.create_index(index_name, table, ['account_id'])


def downgrade() -> None:
    # Revert to nullable
    for table in ('users', 'trees', 'tree_categories', 'tree_tags',
                  'step_categories', 'step_library', 'tree_embeddings', 'feedback'):
        op.alter_column(table, 'account_id', nullable=True)
    for table, index_name in (
        ('tree_embeddings', 'ix_tree_embeddings_account_id'),
        ('feedback', 'ix_feedback_account_id'),
    ):
        try:
            op.drop_index(index_name, table_name=table)
        except Exception:
            pass
