resolutionflow/backend/alembic/versions/018_migrate_users_to_accounts.py

"""migrate existing users and teams to accounts

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

This is the most critical migration. It creates a _team_account_mapping table
for deterministic cross-migration lookups, then migrates all users to accounts.

Three paths:
  A) Teams with users → Account with deterministic owner
  B) Teams with zero users → Account with owner_id=NULL, subscription status='orphaned'
  C) Users without a team → Personal account, user is owner
"""
import secrets
import string
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 = '018'
down_revision: Union[str, None] = '017'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

# Characters for display codes — exclude confusing chars
DISPLAY_CODE_CHARS = string.ascii_uppercase + string.digits
DISPLAY_CODE_CHARS = DISPLAY_CODE_CHARS.replace('0', '').replace('O', '').replace('I', '').replace('1', '').replace('L', '')


def _generate_display_code(existing_codes: set) -> str:
    """Generate a unique 8-character display code."""
    for _ in range(100):
        code = ''.join(secrets.choice(DISPLAY_CODE_CHARS) for _ in range(8))
        if code not in existing_codes:
            existing_codes.add(code)
            return code
    raise RuntimeError("Failed to generate unique display code after 100 attempts")


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

    # Create mapping table for deterministic cross-migration lookups
    op.create_table(
        '_team_account_mapping',
        sa.Column('team_id', UUID(as_uuid=True), nullable=False),
        sa.Column('account_id', UUID(as_uuid=True), nullable=False),
        sa.Column('owner_user_id', UUID(as_uuid=True), nullable=True),
        sa.PrimaryKeyConstraint('team_id'),
    )

    existing_codes: set = set()

    # --- Path A & B: Process all teams ---
    teams = conn.execute(sa.text("SELECT id, name FROM teams")).fetchall()

    for team in teams:
        team_id = team[0]
        team_name = team[1]
        display_code = _generate_display_code(existing_codes)

        # Find deterministic owner: team admin first, then earliest user
        owner_row = conn.execute(sa.text("""
            SELECT id FROM users
            WHERE team_id = :tid
            ORDER BY is_team_admin DESC, created_at ASC, id ASC
            LIMIT 1
        """), {"tid": team_id}).fetchone()

        owner_user_id = owner_row[0] if owner_row else None

        # Create account
        conn.execute(sa.text("""
            INSERT INTO accounts (id, name, display_code, owner_id, created_at, updated_at)
            VALUES (gen_random_uuid(), :name, :code, :owner_id, NOW(), NOW())
        """), {"name": team_name, "code": display_code, "owner_id": owner_user_id})

        # Get the account we just created
        account_row = conn.execute(sa.text(
            "SELECT id FROM accounts WHERE display_code = :code"
        ), {"code": display_code}).fetchone()
        account_id = account_row[0]

        # Insert mapping
        conn.execute(sa.text("""
            INSERT INTO _team_account_mapping (team_id, account_id, owner_user_id)
            VALUES (:tid, :aid, :uid)
        """), {"tid": team_id, "aid": account_id, "uid": owner_user_id})

        if owner_user_id is not None:
            # Path A: Team with users
            # Create active subscription
            conn.execute(sa.text("""
                INSERT INTO subscriptions (id, account_id, plan, status, created_at, updated_at)
                VALUES (gen_random_uuid(), :aid, 'free', 'active', NOW(), NOW())
            """), {"aid": account_id})

            # Update all users in this team
            # Team admins become owners, others keep their role
            conn.execute(sa.text("""
                UPDATE users SET
                    account_id = :aid,
                    account_role = CASE
                        WHEN is_team_admin = true THEN 'owner'
                        ELSE role
                    END
                WHERE team_id = :tid
            """), {"aid": account_id, "tid": team_id})
        else:
            # Path B: Team with zero users (orphan)
            conn.execute(sa.text("""
                INSERT INTO subscriptions (id, account_id, plan, status, created_at, updated_at)
                VALUES (gen_random_uuid(), :aid, 'free', 'orphaned', NOW(), NOW())
            """), {"aid": account_id})

    # --- Path C: Users without a team ---
    teamless_users = conn.execute(sa.text(
        "SELECT id, name FROM users WHERE team_id IS NULL AND account_id IS NULL"
    )).fetchall()

    for user in teamless_users:
        user_id = user[0]
        user_name = user[1]
        display_code = _generate_display_code(existing_codes)

        # Create personal account (owner_id set to NULL initially)
        conn.execute(sa.text("""
            INSERT INTO accounts (id, name, display_code, owner_id, created_at, updated_at)
            VALUES (gen_random_uuid(), :name, :code, NULL, NOW(), NOW())
        """), {"name": f"{user_name}'s Account", "code": display_code})

        account_row = conn.execute(sa.text(
            "SELECT id FROM accounts WHERE display_code = :code"
        ), {"code": display_code}).fetchone()
        account_id = account_row[0]

        # Update user
        conn.execute(sa.text("""
            UPDATE users SET account_id = :aid, account_role = 'owner'
            WHERE id = :uid
        """), {"aid": account_id, "uid": user_id})

        # Set owner
        conn.execute(sa.text("""
            UPDATE accounts SET owner_id = :uid WHERE id = :aid
        """), {"uid": user_id, "aid": account_id})

        # Create free subscription
        conn.execute(sa.text("""
            INSERT INTO subscriptions (id, account_id, plan, status, created_at, updated_at)
            VALUES (gen_random_uuid(), :aid, 'free', 'active', NOW(), NOW())
        """), {"aid": account_id})

    # --- Validation ---
    orphaned_users = conn.execute(sa.text(
        "SELECT COUNT(*) FROM users WHERE account_id IS NULL"
    )).scalar()
    if orphaned_users > 0:
        raise RuntimeError(
            f"Migration 018 failed validation: {orphaned_users} users still have NULL account_id"
        )

    team_count = conn.execute(sa.text("SELECT COUNT(*) FROM teams")).scalar()
    mapping_count = conn.execute(sa.text("SELECT COUNT(*) FROM _team_account_mapping")).scalar()
    if mapping_count != team_count:
        raise RuntimeError(
            f"Migration 018 failed: mapping count ({mapping_count}) != team count ({team_count})"
        )


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

    # Clear account data from users
    conn.execute(sa.text("UPDATE users SET account_id = NULL, account_role = NULL"))

    # Delete all subscriptions and accounts created by this migration
    conn.execute(sa.text("DELETE FROM subscriptions"))
    conn.execute(sa.text("DELETE FROM accounts"))

    # Drop mapping table
    op.drop_table('_team_account_mapping')