"""Create traffic_data table for storing traffic data for re-training

Revision ID: 001_traffic_data
Revises: 
Create Date: 2025-01-08 12:00:00.000000

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

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


def upgrade():
    """Create traffic_data table"""
    op.create_table('traffic_data',
        sa.Column('id', UUID(as_uuid=True), nullable=False, primary_key=True),
        sa.Column('location_id', sa.String(100), nullable=False, index=True),
        sa.Column('date', sa.DateTime(timezone=True), nullable=False, index=True),
        sa.Column('traffic_volume', sa.Integer, nullable=True),
        sa.Column('pedestrian_count', sa.Integer, nullable=True),
        sa.Column('congestion_level', sa.String(20), nullable=True),
        sa.Column('average_speed', sa.Float, nullable=True),
        sa.Column('source', sa.String(50), nullable=False, server_default='madrid_opendata'),
        sa.Column('raw_data', sa.Text, nullable=True),
        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False),
        sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False),
    )
    
    # Create index for efficient querying by location and date
    op.create_index(
        'idx_traffic_location_date', 
        'traffic_data', 
        ['location_id', 'date']
    )
    
    # Create index for date range queries  
    op.create_index(
        'idx_traffic_date_range',
        'traffic_data',
        ['date']
    )


def downgrade():
    """Drop traffic_data table"""
    op.drop_index('idx_traffic_date_range', table_name='traffic_data')
    op.drop_index('idx_traffic_location_date', table_name='traffic_data')
    op.drop_table('traffic_data')