"""Add asset_groups and policies

Revision ID: 005
Revises: 004
Create Date: 2026-02-01 11:45:00.000000

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '005'
down_revision = '004'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # 1. Create association table asset_groups
    op.create_table(
        'asset_groups',
        sa.Column('asset_id', sa.Integer(), nullable=False),
        sa.Column('group_id', sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(['asset_id'], ['assets.id'], ),
        sa.ForeignKeyConstraint(['group_id'], ['groups.id'], ),
        sa.PrimaryKeyConstraint('asset_id', 'group_id')
    )

    # 2. Data Migration: Copy existing assigned_group_id to new table
    # Only if assigned_group_id exists (check logic handled by migration existing)
    op.execute(
        "INSERT INTO asset_groups (asset_id, group_id) "
        "SELECT id, assigned_group_id FROM assets "
        "WHERE assigned_group_id IS NOT NULL"
    )

    # 3. Drop old relationship on assets
    # Note: We need to drop the FK constraint first.
    # Name explicitly given in 002_combined_updates.py: 'fk_assets_group'
    op.drop_constraint('fk_assets_group', 'assets', type_='foreignkey')
    op.drop_column('assets', 'assigned_group_id')

    # 4. Add policy_id to groups
    op.add_column('groups', sa.Column('policy_id', sa.Integer(), nullable=True))
    op.create_foreign_key('fk_groups_policy_id', 'groups', 'policies', ['policy_id'], ['id'])


def downgrade() -> None:
    # Reverse order
    
    # 1. Drop policy_id from groups
    op.drop_constraint('fk_groups_policy_id', 'groups', type_='foreignkey')
    op.drop_column('groups', 'policy_id')

    # 2. Restore assigned_group_id to assets
    op.add_column('assets', sa.Column('assigned_group_id', sa.Integer(), nullable=True))
    op.create_foreign_key('fk_assets_group', 'assets', 'groups', ['assigned_group_id'], ['id'])

    # 3. Data Migration: Restore (optional/partial) - picking one group if multiple
    op.execute(
        "UPDATE assets SET assigned_group_id = ("
        "  SELECT group_id FROM asset_groups WHERE asset_groups.asset_id = assets.id LIMIT 1"
        ")"
    )

    # 4. Drop asset_groups table
    op.drop_table('asset_groups')
