"""Initial coverage tables

Revision ID: 0001
Revises: 
Create Date: 2024-09-14 12:00:00.000000

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

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


def upgrade() -> None:
    # Create coverage_versions table
    op.create_table(
        'coverage_versions',
        sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
        sa.Column('version', sa.String(length=50), nullable=False),
        sa.Column('jurisdiction', sa.String(length=10), nullable=False),
        sa.Column('tax_year', sa.String(length=10), nullable=False),
        sa.Column('tenant_id', sa.String(length=100), nullable=True),
        sa.Column('source_files', postgresql.JSON(astext_type=sa.Text()), nullable=False),
        sa.Column('compiled_at', sa.DateTime(), nullable=False),
        sa.Column('hash', sa.String(length=64), nullable=False),
        sa.PrimaryKeyConstraint('id')
    )
    
    # Create indexes for coverage_versions
    op.create_index('ix_coverage_versions_version', 'coverage_versions', ['version'])
    op.create_index('ix_coverage_versions_jurisdiction_tax_year', 'coverage_versions', ['jurisdiction', 'tax_year'])
    op.create_index('ix_coverage_versions_tenant_id', 'coverage_versions', ['tenant_id'])
    op.create_index('ix_coverage_versions_hash', 'coverage_versions', ['hash'])
    
    # Create coverage_audit table
    op.create_table(
        'coverage_audit',
        sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
        sa.Column('taxpayer_id', sa.String(length=100), nullable=False),
        sa.Column('tax_year', sa.String(length=10), nullable=False),
        sa.Column('policy_version', sa.String(length=50), nullable=False),
        sa.Column('overall_status', sa.String(length=20), nullable=False),
        sa.Column('blocking_items', postgresql.JSON(astext_type=sa.Text()), nullable=False),
        sa.Column('created_at', sa.DateTime(), nullable=False),
        sa.Column('trace_id', sa.String(length=100), nullable=True),
        sa.PrimaryKeyConstraint('id')
    )
    
    # Create indexes for coverage_audit
    op.create_index('ix_coverage_audit_taxpayer_id', 'coverage_audit', ['taxpayer_id'])
    op.create_index('ix_coverage_audit_tax_year', 'coverage_audit', ['tax_year'])
    op.create_index('ix_coverage_audit_taxpayer_tax_year', 'coverage_audit', ['taxpayer_id', 'tax_year'])
    op.create_index('ix_coverage_audit_created_at', 'coverage_audit', ['created_at'])
    op.create_index('ix_coverage_audit_trace_id', 'coverage_audit', ['trace_id'])


def downgrade() -> None:
    # Drop coverage_audit table and indexes
    op.drop_index('ix_coverage_audit_trace_id', table_name='coverage_audit')
    op.drop_index('ix_coverage_audit_created_at', table_name='coverage_audit')
    op.drop_index('ix_coverage_audit_taxpayer_tax_year', table_name='coverage_audit')
    op.drop_index('ix_coverage_audit_tax_year', table_name='coverage_audit')
    op.drop_index('ix_coverage_audit_taxpayer_id', table_name='coverage_audit')
    op.drop_table('coverage_audit')
    
    # Drop coverage_versions table and indexes
    op.drop_index('ix_coverage_versions_hash', table_name='coverage_versions')
    op.drop_index('ix_coverage_versions_tenant_id', table_name='coverage_versions')
    op.drop_index('ix_coverage_versions_jurisdiction_tax_year', table_name='coverage_versions')
    op.drop_index('ix_coverage_versions_version', table_name='coverage_versions')
    op.drop_table('coverage_versions')