import sqlalchemy as sa
from alembic import op

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

def upgrade():
    op.create_table('data_sources', sa.Column('id', sa.Integer(), nullable=False), sa.Column('code', sa.String(50), nullable=False, unique=True), sa.Column('name', sa.String(200), nullable=False), sa.Column('source_type', sa.String(50), nullable=True), sa.Column('region_level', sa.String(20), nullable=True), sa.Column('region_code', sa.String(20), nullable=True), sa.Column('adapter_class', sa.String(200), nullable=True), sa.Column('adapter_config', sa.JSON(), nullable=True), sa.Column('has_categories', sa.Boolean(), nullable=False, server_default=sa.false()), sa.Column('crawl_schedule', sa.String(100), nullable=True), sa.Column('request_delay_min', sa.Float(), nullable=False, server_default='1.0'), sa.Column('request_delay_max', sa.Float(), nullable=False, server_default='5.0'), sa.Column('max_retries', sa.Integer(), nullable=False, server_default='3'), sa.Column('is_active', sa.Boolean(), nullable=False, server_default=sa.true()), sa.Column('last_crawled_at', sa.DateTime(), nullable=True), sa.Column('created_at', sa.DateTime(), nullable=True), sa.PrimaryKeyConstraint('id'), schema='data_center')
    op.execute("INSERT INTO data_center.data_sources (id, code, name, source_type, region_level, region_code, adapter_class, has_categories, is_active) VALUES (1, 'chinatax', '国家税务总局', 'regulation', 'national', 'CN', 'adapters.chinatax.adapter.ChinataxAdapter', TRUE, TRUE)")
    op.add_column('tax_documents', sa.Column('source_id', sa.Integer(), sa.ForeignKey('data_center.data_sources.id'), nullable=True), schema='data_center')
    op.add_column('tax_documents', sa.Column('doc_type', sa.String(20), nullable=True), schema='data_center')
    op.add_column('tax_documents', sa.Column('region_code', sa.String(20), nullable=True), schema='data_center')
    op.add_column('tax_documents', sa.Column('qa_question', sa.Text(), nullable=True), schema='data_center')
    op.add_column('tax_documents', sa.Column('qa_answer', sa.Text(), nullable=True), schema='data_center')
    op.execute("UPDATE data_center.tax_documents SET source_id=1, doc_type='regulation', region_code='CN' WHERE source_id IS NULL")
    op.create_index('idx_doc_source_id', 'tax_documents', ['source_id'], schema='data_center')
    op.create_index('idx_doc_doc_type', 'tax_documents', ['doc_type'], schema='data_center')
    op.create_index('idx_doc_region_code', 'tax_documents', ['region_code'], schema='data_center')
    op.add_column('data_processing_tasks', sa.Column('source_id', sa.Integer(), sa.ForeignKey('data_center.data_sources.id'), nullable=True), schema='data_center')

def downgrade():
    op.drop_column('data_processing_tasks', 'source_id', schema='data_center')
    op.drop_index('idx_doc_region_code', table_name='tax_documents', schema='data_center')
    op.drop_index('idx_doc_doc_type', table_name='tax_documents', schema='data_center')
    op.drop_index('idx_doc_source_id', table_name='tax_documents', schema='data_center')
    op.drop_column('tax_documents', 'qa_answer', schema='data_center')
    op.drop_column('tax_documents', 'qa_question', schema='data_center')
    op.drop_column('tax_documents', 'region_code', schema='data_center')
    op.drop_column('tax_documents', 'doc_type', schema='data_center')
    op.drop_column('tax_documents', 'source_id', schema='data_center')
    op.drop_table('data_sources', schema='data_center')