"""Add EPSS and CISA KEV enrichment fields to vulnerabilities

Revision ID: 007
Revises: 006
Create Date: 2026-05-11 12:00:00.000000

Adds columns for:
- EPSS (Exploit Prediction Scoring System) from FIRST.org
- CISA KEV (Known Exploited Vulnerabilities) catalog
- Generic enrichment tracking (sources used, last updated)
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


revision = '007'
down_revision = '006'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # EPSS fields
    op.add_column(
        'vulnerabilities',
        sa.Column('epss_score', sa.Float(), nullable=True)
    )
    op.add_column(
        'vulnerabilities',
        sa.Column('epss_percentile', sa.Float(), nullable=True)
    )
    op.add_column(
        'vulnerabilities',
        sa.Column('epss_updated_at', sa.DateTime(), nullable=True)
    )
    op.create_index('ix_vulnerabilities_epss_score', 'vulnerabilities', ['epss_score'])

    # CISA KEV fields
    op.add_column(
        'vulnerabilities',
        sa.Column('kev_listed', sa.Boolean(), server_default='false', nullable=False)
    )
    op.add_column(
        'vulnerabilities',
        sa.Column('kev_date_added', sa.DateTime(), nullable=True)
    )
    op.add_column(
        'vulnerabilities',
        sa.Column('kev_ransomware_use', sa.Boolean(), server_default='false', nullable=False)
    )
    op.add_column(
        'vulnerabilities',
        sa.Column('kev_short_description', sa.Text(), nullable=True)
    )
    op.create_index('ix_vulnerabilities_kev_listed', 'vulnerabilities', ['kev_listed'])

    # Generic enrichment tracking
    op.add_column(
        'vulnerabilities',
        sa.Column('enrichment_sources', sa.Text(), nullable=True)
    )
    op.add_column(
        'vulnerabilities',
        sa.Column('enrichment_updated_at', sa.DateTime(), nullable=True)
    )


def downgrade() -> None:
    op.drop_index('ix_vulnerabilities_kev_listed', table_name='vulnerabilities')
    op.drop_index('ix_vulnerabilities_epss_score', table_name='vulnerabilities')

    op.drop_column('vulnerabilities', 'enrichment_updated_at')
    op.drop_column('vulnerabilities', 'enrichment_sources')
    op.drop_column('vulnerabilities', 'kev_short_description')
    op.drop_column('vulnerabilities', 'kev_ransomware_use')
    op.drop_column('vulnerabilities', 'kev_date_added')
    op.drop_column('vulnerabilities', 'kev_listed')
    op.drop_column('vulnerabilities', 'epss_updated_at')
    op.drop_column('vulnerabilities', 'epss_percentile')
    op.drop_column('vulnerabilities', 'epss_score')
