"""Combined migration matching server state

Revision ID: 002
Revises: 001
Create Date: 2026-01-30 14:15:00.000000

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '002'
down_revision = '001'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # --- Logic from original 002 ---
    # 1. Erstelle Enums
    # (Alembic erstellt Enums automatisch bei der ersten Nutzung in create_table, 
    # aber wir brauchen sie für die Spalten unten)
    policy_status = sa.Enum('ACTIVE', 'WARNING', 'DISABLED', name='policystatus')
    notif_type = sa.Enum('SLA_BREACH', 'ASSIGNMENT', 'NEW_VULNERABILITY', 'MANUAL', name='notificationtype')
    notif_status = sa.Enum('SENT', 'FAILED', 'SUPPRESSED', name='notificationstatus')
    interval_type = sa.Enum('EVERY_HOUR', 'EVERY_6_HOURS', 'EVERY_12_HOURS', 'DAILY', 'WEEKLY', 'CRON', name='scheduleinterval')

    # 2. Erstelle neue Tabellen
    op.create_table(
        'groups',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=100), nullable=False),
        sa.Column('description', sa.Text(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.Column('updated_at', sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('name')
    )
    op.create_index('ix_groups_id', 'groups', ['id'])

    op.create_table(
        'user_groups',
        sa.Column('user_id', sa.Integer(), nullable=False),
        sa.Column('group_id', sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(['user_id'], ['users.id']),
        sa.ForeignKeyConstraint(['group_id'], ['groups.id']),
        sa.PrimaryKeyConstraint('user_id', 'group_id')
    )

    op.create_table(
        'policies',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=255), nullable=False),
        sa.Column('description', sa.Text(), nullable=True),
        sa.Column('category', sa.String(length=100), nullable=False),
        sa.Column('status', policy_status, nullable=True),
        sa.Column('critical_sla_days', sa.Integer(), nullable=True),
        sa.Column('high_sla_days', sa.Integer(), nullable=True),
        sa.Column('medium_sla_days', sa.Integer(), nullable=True),
        sa.Column('low_sla_days', sa.Integer(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.Column('updated_at', sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint('id')
    )
    op.create_index('ix_policies_id', 'policies', ['id'])

    op.create_table(
        'scan_schedules',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=200), nullable=False),
        sa.Column('interval', interval_type, nullable=False),
        sa.Column('cron_expression', sa.String(length=100), nullable=True),
        sa.Column('enabled', sa.Boolean(), nullable=False),
        sa.Column('last_run', sa.DateTime(), nullable=True),
        sa.Column('next_run', sa.DateTime(), nullable=True),
        sa.Column('created_by', sa.Integer(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.Column('updated_at', sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint('id')
    )
    op.create_index('ix_scan_schedules_id', 'scan_schedules', ['id'])

    op.create_table(
        'notification_logs',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('vulnerability_id', sa.Integer(), nullable=True),
        sa.Column('asset_id', sa.Integer(), nullable=True),
        sa.Column('user_id', sa.Integer(), nullable=True),
        sa.Column('notification_type', notif_type, nullable=False),
        sa.Column('sent_at', sa.DateTime(), nullable=False),
        sa.Column('subject', sa.String(length=500), nullable=True),
        sa.Column('recipient_email', sa.String(length=255), nullable=True),
        sa.Column('status', notif_status, nullable=False),
        sa.Column('message_body', sa.Text(), nullable=True),
        sa.Column('error_message', sa.Text(), nullable=True),
        sa.ForeignKeyConstraint(['vulnerability_id'], ['vulnerabilities.id']),
        sa.ForeignKeyConstraint(['asset_id'], ['assets.id']),
        sa.ForeignKeyConstraint(['user_id'], ['users.id']),
        sa.PrimaryKeyConstraint('id')
    )
    op.create_table(
        'settings',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('key', sa.String(), nullable=False),
        sa.Column('value', sa.Text(), nullable=True),
        sa.Column('description', sa.String(), nullable=True),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.Column('updated_at', sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint('id')
    )
    op.create_index('ix_settings_id', 'settings', ['id'])
    op.create_index('ix_settings_key', 'settings', ['key'], unique=True)

    # 3. Fehlende Spalten zu Assets hinzufügen
    op.add_column('assets', sa.Column('assigned_user_id', sa.Integer(), nullable=True))
    op.add_column('assets', sa.Column('assigned_group_id', sa.Integer(), nullable=True))
    op.add_column('assets', sa.Column('policy_id', sa.Integer(), nullable=True))
    op.create_foreign_key('fk_assets_user', 'assets', 'users', ['assigned_user_id'], ['id'])
    op.create_foreign_key('fk_assets_group', 'assets', 'groups', ['assigned_group_id'], ['id'])
    op.create_foreign_key('fk_assets_policy', 'assets', 'policies', ['policy_id'], ['id'])
    op.create_index('ix_assets_assigned_user_id', 'assets', ['assigned_user_id'])
    op.create_index('ix_assets_assigned_group_id', 'assets', ['assigned_group_id'])

    # 4. Fehlende Spalten zu Vulnerabilities hinzufügen
    op.add_column('vulnerabilities', sa.Column('assigned_user_id', sa.Integer(), nullable=True))
    op.add_column('vulnerabilities', sa.Column('assigned_group_id', sa.Integer(), nullable=True))
    op.add_column('vulnerabilities', sa.Column('deferred_until', sa.DateTime(), nullable=True))
    # sa.Column('defer_reason', sa.Text(), nullable=True), # Nur Info
    op.add_column('vulnerabilities', sa.Column('defer_reason', sa.Text(), nullable=True))
    op.add_column('vulnerabilities', sa.Column('notification_suppressed', sa.Boolean(), server_default='false', nullable=False))
    op.add_column('vulnerabilities', sa.Column('priority_score', sa.Float(), server_default='0', nullable=True))
    op.add_column('vulnerabilities', sa.Column('sla_breach_date', sa.DateTime(), nullable=True))
    
    op.create_foreign_key('fk_vulnerabilities_user', 'vulnerabilities', 'users', ['assigned_user_id'], ['id'])
    op.create_foreign_key('fk_vulnerabilities_group', 'vulnerabilities', 'groups', ['assigned_group_id'], ['id'])
    op.create_index('ix_vulnerabilities_assigned_user_id', 'vulnerabilities', ['assigned_user_id'])
    op.create_index('ix_vulnerabilities_assigned_group_id', 'vulnerabilities', ['assigned_group_id'])

    # --- Logic from original enum fix ---
    # Add new enum value to notificationtype if it doesn't exist.
    op.execute("ALTER TYPE notificationtype ADD VALUE IF NOT EXISTS 'NEW_VULNERABILITY'")


def downgrade() -> None:
    # --- Downgrade Logic ---
    op.drop_table('notification_logs')
    op.drop_table('scan_schedules')
    op.drop_table('policies')
    op.drop_table('user_groups')
    op.drop_table('groups')
    
    op.drop_column('assets', 'assigned_user_id')
    op.drop_column('assets', 'assigned_group_id')
    op.drop_column('assets', 'policy_id')
    
    op.drop_column('vulnerabilities', 'assigned_user_id')
    op.drop_column('vulnerabilities', 'assigned_group_id')
    op.drop_column('vulnerabilities', 'deferred_until')
    op.drop_column('vulnerabilities', 'defer_reason')
    op.drop_column('vulnerabilities', 'notification_suppressed')
    op.drop_column('vulnerabilities', 'priority_score')
    op.drop_column('vulnerabilities', 'sla_breach_date')

    op.execute("DROP TYPE IF EXISTS policystatus")
    op.execute("DROP TYPE IF EXISTS notificationtype")
    op.execute("DROP TYPE IF EXISTS notificationstatus")
    op.execute("DROP TYPE IF EXISTS scheduleinterval")
