"""Add scan progress tracking

Revision ID: 012
Revises: 011
Create Date: 2024-01-01 00:00:00.000000

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '012'
down_revision = '011'
branch_labels = None
depends_on = None


def upgrade():
    # Add progress tracking columns to scans table
    op.add_column('scans', sa.Column('current_phase', sa.String(50), nullable=True,
                  comment='Current scan phase: ping, tcp_scan, udp_scan, service_detection, http_analysis'))
    op.add_column('scans', sa.Column('total_ips', sa.Integer(), nullable=True,
                  comment='Total number of IPs to scan'))
    op.add_column('scans', sa.Column('completed_ips', sa.Integer(), nullable=True, default=0,
                  comment='Number of IPs completed in current phase'))

    # Create scan_progress table for per-IP progress tracking
    op.create_table(
        'scan_progress',
        sa.Column('id', sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column('scan_id', sa.Integer(), sa.ForeignKey('scans.id'), nullable=False, index=True),
        sa.Column('ip_address', sa.String(45), nullable=False, comment='IP address being scanned'),
        sa.Column('site_name', sa.String(255), nullable=True, comment='Site name this IP belongs to'),
        sa.Column('phase', sa.String(50), nullable=False,
                  comment='Phase: ping, tcp_scan, udp_scan, service_detection, http_analysis'),
        sa.Column('status', sa.String(20), nullable=False, default='pending',
                  comment='pending, in_progress, completed, failed'),
        sa.Column('ping_result', sa.Boolean(), nullable=True, comment='Ping response result'),
        sa.Column('tcp_ports', sa.Text(), nullable=True, comment='JSON array of discovered TCP ports'),
        sa.Column('udp_ports', sa.Text(), nullable=True, comment='JSON array of discovered UDP ports'),
        sa.Column('services', sa.Text(), nullable=True, comment='JSON array of detected services'),
        sa.Column('created_at', sa.DateTime(), nullable=False, server_default=sa.func.now(),
                  comment='Entry creation time'),
        sa.Column('updated_at', sa.DateTime(), nullable=False, server_default=sa.func.now(),
                  onupdate=sa.func.now(), comment='Last update time'),
        sa.UniqueConstraint('scan_id', 'ip_address', name='uix_scan_progress_ip')
    )


def downgrade():
    # Drop scan_progress table
    op.drop_table('scan_progress')

    # Remove progress tracking columns from scans table
    op.drop_column('scans', 'completed_ips')
    op.drop_column('scans', 'total_ips')
    op.drop_column('scans', 'current_phase')