"""Consolidate user and superadmin sessions into unified sessions table.

Revision ID: c8d2e4f6a1b3
Revises: b7e3f1a92c4d
Create Date: 2026-04-28 00:00:00.000000

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

# revision identifiers, used by Alembic.
revision = 'c8d2e4f6a1b3'
down_revision = 'b7e3f1a92c4d'
branch_labels = None
depends_on = None


def upgrade():
    # 1. Add new columns (nullable initially for data migration)
    op.add_column('sessions', sa.Column('owner_type', sa.String(20), nullable=True))
    op.add_column('sessions', sa.Column('owner_id', sa.String(36), nullable=True))

    # 2. Backfill existing user sessions: owner_type = 'user', owner_id = user_id
    op.execute("""
        UPDATE sessions
        SET owner_type = 'user',
            owner_id = user_id
        WHERE owner_type IS NULL
    """)

    # 3. Migrate superadmin sessions into the sessions table
    op.execute("""
        INSERT INTO sessions (
            id, owner_type, owner_id, token, status,
            ip_address, user_agent, device_info,
            expires_at, last_activity_at, revoked_at, revoked_reason,
            is_compliance_only, created_at, updated_at, deleted_at
        )
        SELECT
            id, 'superadmin', superadmin_id, token, 'active',
            ip_address, user_agent, NULL,
            expires_at, last_activity_at, revoked_at, revoked_reason,
            FALSE, created_at, updated_at, deleted_at
        FROM superadmin_sessions
    """)

    # 4. Make owner_type and owner_id NOT NULL
    op.alter_column('sessions', 'owner_type', nullable=False)
    op.alter_column('sessions', 'owner_id', nullable=False)

    # 5. Make user_id nullable (no longer the sole owner reference)
    op.alter_column('sessions', 'user_id', nullable=True)

    # 6. Create indexes for efficient owner-scoped queries
    op.create_index(
        'ix_sessions_owner_type_owner_id',
        'sessions',
        ['owner_type', 'owner_id']
    )
    op.create_index(
        'ix_sessions_owner_type',
        'sessions',
        ['owner_type']
    )
    op.create_index(
        'ix_sessions_owner_id',
        'sessions',
        ['owner_id']
    )

    # 7. Drop the now-redundant superadmin_sessions table
    op.drop_table('superadmin_sessions')


def downgrade():
    # 1. Recreate superadmin_sessions table
    op.create_table(
        'superadmin_sessions',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('superadmin_id', sa.String(36), sa.ForeignKey('superadmins.id'), nullable=False, index=True),
        sa.Column('token', sa.String(255), unique=True, nullable=False, index=True),
        sa.Column('expires_at', sa.DateTime, nullable=False),
        sa.Column('last_activity_at', sa.DateTime, nullable=False),
        sa.Column('ip_address', sa.String(45), nullable=True),
        sa.Column('user_agent', sa.Text, nullable=True),
        sa.Column('revoked_at', sa.DateTime, nullable=True),
        sa.Column('revoked_reason', sa.String(255), nullable=True),
        sa.Column('created_at', sa.DateTime, nullable=False),
        sa.Column('updated_at', sa.DateTime, nullable=False),
        sa.Column('deleted_at', sa.DateTime, nullable=True),
    )

    # 2. Move superadmin sessions back to superadmin_sessions
    op.execute("""
        INSERT INTO superadmin_sessions (
            id, superadmin_id, token, expires_at, last_activity_at,
            ip_address, user_agent, revoked_at, revoked_reason,
            created_at, updated_at, deleted_at
        )
        SELECT
            id, owner_id, token, expires_at, last_activity_at,
            ip_address, user_agent, revoked_at, revoked_reason,
            created_at, updated_at, deleted_at
        FROM sessions
        WHERE owner_type = 'superadmin'
    """)

    # 3. Remove superadmin sessions from sessions table
    op.execute("DELETE FROM sessions WHERE owner_type = 'superadmin'")

    # 4. Drop indexes
    op.drop_index('ix_sessions_owner_id', table_name='sessions')
    op.drop_index('ix_sessions_owner_type', table_name='sessions')
    op.drop_index('ix_sessions_owner_type_owner_id', table_name='sessions')

    # 5. Remove new columns
    op.drop_column('sessions', 'owner_id')
    op.drop_column('sessions', 'owner_type')

    # 6. Make user_id NOT NULL again
    op.alter_column('sessions', 'user_id', nullable=False)