"""finalize account migration — add constraints, clean up orphans

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

Adds NOT NULL constraints, foreign keys, and CHECK constraints.
Cleans up orphan accounts (zero-user teams with no content).
"""
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 = '020'
down_revision: Union[str, None] = '019'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

CONTENT_TABLES = ['trees', 'step_library', 'tree_categories', 'tree_tags', 'step_categories']


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

    # 1. Clean up orphan accounts (zero-user teams with no content)
    conn.execute(sa.text("""
        DELETE FROM subscriptions WHERE account_id IN (
            SELECT a.id FROM accounts a
            WHERE a.owner_id IS NULL
            AND NOT EXISTS (SELECT 1 FROM trees t WHERE t.account_id = a.id)
            AND NOT EXISTS (SELECT 1 FROM tree_categories tc WHERE tc.account_id = a.id)
            AND NOT EXISTS (SELECT 1 FROM tree_tags tt WHERE tt.account_id = a.id)
            AND NOT EXISTS (SELECT 1 FROM step_categories sc WHERE sc.account_id = a.id)
            AND NOT EXISTS (SELECT 1 FROM step_library sl WHERE sl.account_id = a.id)
        )
    """))

    # Also remove the mapping entries for these orphans
    conn.execute(sa.text("""
        DELETE FROM _team_account_mapping WHERE account_id IN (
            SELECT a.id FROM accounts a
            WHERE a.owner_id IS NULL
            AND NOT EXISTS (SELECT 1 FROM trees t WHERE t.account_id = a.id)
            AND NOT EXISTS (SELECT 1 FROM tree_categories tc WHERE tc.account_id = a.id)
            AND NOT EXISTS (SELECT 1 FROM tree_tags tt WHERE tt.account_id = a.id)
            AND NOT EXISTS (SELECT 1 FROM step_categories sc WHERE sc.account_id = a.id)
            AND NOT EXISTS (SELECT 1 FROM step_library sl WHERE sl.account_id = a.id)
        )
    """))

    conn.execute(sa.text("""
        DELETE FROM accounts
        WHERE owner_id IS NULL
        AND NOT EXISTS (SELECT 1 FROM trees t WHERE t.account_id = accounts.id)
        AND NOT EXISTS (SELECT 1 FROM tree_categories tc WHERE tc.account_id = accounts.id)
        AND NOT EXISTS (SELECT 1 FROM tree_tags tt WHERE tt.account_id = accounts.id)
        AND NOT EXISTS (SELECT 1 FROM step_categories sc WHERE sc.account_id = accounts.id)
        AND NOT EXISTS (SELECT 1 FROM step_library sl WHERE sl.account_id = accounts.id)
    """))

    # 2. Users: enforce NOT NULL and add FK + CHECK
    op.alter_column('users', 'account_id', nullable=False)
    op.alter_column('users', 'account_role', nullable=False)
    op.create_foreign_key(
        'fk_users_account_id', 'users', 'accounts',
        ['account_id'], ['id'], ondelete='CASCADE'
    )
    op.create_check_constraint(
        'ck_users_account_role_enum', 'users',
        "account_role IN ('owner', 'engineer', 'viewer')"
    )

    # 3. Content tables: add FK on account_id (nullable OK — NULL means global)
    for table in CONTENT_TABLES:
        op.create_foreign_key(
            f'fk_{table}_account_id', table, 'accounts',
            ['account_id'], ['id'], ondelete='CASCADE'
        )

    # 4. Accounts: add owner FK (owner_id stays nullable due to circular FK with users)
    op.create_foreign_key(
        'fk_accounts_owner_id', 'accounts', 'users',
        ['owner_id'], ['id'], ondelete='RESTRICT'
    )


def downgrade() -> None:
    # Remove account owner FK and nullable constraint
    op.drop_constraint('fk_accounts_owner_id', 'accounts', type_='foreignkey')
    op.alter_column('accounts', 'owner_id', nullable=True)

    # Remove content table FKs
    for table in reversed(CONTENT_TABLES):
        op.drop_constraint(f'fk_{table}_account_id', table, type_='foreignkey')

    # Remove user constraints
    op.drop_constraint('ck_users_account_role_enum', 'users', type_='check')
    op.drop_constraint('fk_users_account_id', 'users', type_='foreignkey')
    op.alter_column('users', 'account_role', nullable=True)
    op.alter_column('users', 'account_id', nullable=True)