"""Add IoT equipment support

Revision ID: 002_add_iot_equipment_support
Revises: 001_unified_initial_schema
Create Date: 2025-01-12 10:00:00.000000

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

# revision identifiers, used by Alembic.
revision = '002_add_iot_equipment_support'
down_revision = '001_unified_initial_schema'
branch_labels = None
depends_on = None


def upgrade():
    """Add IoT connectivity fields to equipment and create sensor data tables"""

    # Add IoT connectivity fields to equipment table
    op.add_column('equipment', sa.Column('iot_enabled', sa.Boolean(), nullable=False, server_default='false'))
    op.add_column('equipment', sa.Column('iot_protocol', sa.String(50), nullable=True))
    op.add_column('equipment', sa.Column('iot_endpoint', sa.String(500), nullable=True))
    op.add_column('equipment', sa.Column('iot_port', sa.Integer(), nullable=True))
    op.add_column('equipment', sa.Column('iot_credentials', postgresql.JSON(astext_type=sa.Text()), nullable=True))
    op.add_column('equipment', sa.Column('iot_connection_status', sa.String(50), nullable=True))
    op.add_column('equipment', sa.Column('iot_last_connected', sa.DateTime(timezone=True), nullable=True))
    op.add_column('equipment', sa.Column('iot_config', postgresql.JSON(astext_type=sa.Text()), nullable=True))
    op.add_column('equipment', sa.Column('manufacturer', sa.String(100), nullable=True))
    op.add_column('equipment', sa.Column('firmware_version', sa.String(50), nullable=True))

    # Add real-time monitoring fields
    op.add_column('equipment', sa.Column('supports_realtime', sa.Boolean(), nullable=False, server_default='false'))
    op.add_column('equipment', sa.Column('poll_interval_seconds', sa.Integer(), nullable=True))

    # Add sensor capability fields
    op.add_column('equipment', sa.Column('temperature_zones', sa.Integer(), nullable=True))
    op.add_column('equipment', sa.Column('supports_humidity', sa.Boolean(), nullable=False, server_default='false'))
    op.add_column('equipment', sa.Column('supports_energy_monitoring', sa.Boolean(), nullable=False, server_default='false'))
    op.add_column('equipment', sa.Column('supports_remote_control', sa.Boolean(), nullable=False, server_default='false'))

    # Create equipment_sensor_readings table for time-series data
    op.create_table(
        'equipment_sensor_readings',
        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True),
        sa.Column('tenant_id', postgresql.UUID(as_uuid=True), nullable=False, index=True),
        sa.Column('equipment_id', postgresql.UUID(as_uuid=True), nullable=False, index=True),
        sa.Column('batch_id', postgresql.UUID(as_uuid=True), nullable=True, index=True),

        # Timestamp
        sa.Column('reading_time', sa.DateTime(timezone=True), nullable=False, index=True),

        # Temperature readings (support multiple zones)
        sa.Column('temperature', sa.Float(), nullable=True),
        sa.Column('temperature_zones', postgresql.JSON(astext_type=sa.Text()), nullable=True),
        sa.Column('target_temperature', sa.Float(), nullable=True),

        # Humidity
        sa.Column('humidity', sa.Float(), nullable=True),
        sa.Column('target_humidity', sa.Float(), nullable=True),

        # Energy monitoring
        sa.Column('energy_consumption_kwh', sa.Float(), nullable=True),
        sa.Column('power_current_kw', sa.Float(), nullable=True),

        # Equipment status
        sa.Column('operational_status', sa.String(50), nullable=True),
        sa.Column('cycle_stage', sa.String(100), nullable=True),
        sa.Column('cycle_progress_percentage', sa.Float(), nullable=True),
        sa.Column('time_remaining_minutes', sa.Integer(), nullable=True),

        # Process parameters
        sa.Column('motor_speed_rpm', sa.Float(), nullable=True),
        sa.Column('door_status', sa.String(20), nullable=True),
        sa.Column('steam_level', sa.Float(), nullable=True),

        # Quality indicators
        sa.Column('product_weight_kg', sa.Float(), nullable=True),
        sa.Column('moisture_content', sa.Float(), nullable=True),

        # Additional sensor data (flexible JSON for manufacturer-specific metrics)
        sa.Column('additional_sensors', postgresql.JSON(astext_type=sa.Text()), nullable=True),

        # Data quality
        sa.Column('data_quality_score', sa.Float(), nullable=True),
        sa.Column('is_anomaly', sa.Boolean(), nullable=False, server_default='false'),

        # Timestamps
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),

        # Foreign key constraints
        sa.ForeignKeyConstraint(['equipment_id'], ['equipment.id'], ondelete='CASCADE'),
    )

    # Create indexes for time-series queries
    op.create_index(
        'idx_sensor_readings_equipment_time',
        'equipment_sensor_readings',
        ['equipment_id', 'reading_time'],
    )
    op.create_index(
        'idx_sensor_readings_batch',
        'equipment_sensor_readings',
        ['batch_id', 'reading_time'],
    )
    op.create_index(
        'idx_sensor_readings_tenant_time',
        'equipment_sensor_readings',
        ['tenant_id', 'reading_time'],
    )

    # Create equipment_connection_logs table for tracking connectivity
    op.create_table(
        'equipment_connection_logs',
        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True),
        sa.Column('tenant_id', postgresql.UUID(as_uuid=True), nullable=False, index=True),
        sa.Column('equipment_id', postgresql.UUID(as_uuid=True), nullable=False, index=True),

        # Connection event
        sa.Column('event_type', sa.String(50), nullable=False),  # connected, disconnected, error, timeout
        sa.Column('event_time', sa.DateTime(timezone=True), nullable=False, index=True),

        # Connection details
        sa.Column('connection_status', sa.String(50), nullable=False),
        sa.Column('protocol_used', sa.String(50), nullable=True),
        sa.Column('endpoint', sa.String(500), nullable=True),

        # Error tracking
        sa.Column('error_message', sa.Text(), nullable=True),
        sa.Column('error_code', sa.String(50), nullable=True),

        # Performance metrics
        sa.Column('response_time_ms', sa.Integer(), nullable=True),
        sa.Column('data_points_received', sa.Integer(), nullable=True),

        # Additional details
        sa.Column('additional_data', postgresql.JSON(astext_type=sa.Text()), nullable=True),

        # Timestamps
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),

        # Foreign key constraints
        sa.ForeignKeyConstraint(['equipment_id'], ['equipment.id'], ondelete='CASCADE'),
    )

    # Create index for connection logs
    op.create_index(
        'idx_connection_logs_equipment_time',
        'equipment_connection_logs',
        ['equipment_id', 'event_time'],
    )

    # Create equipment_alerts table for IoT-based alerts
    op.create_table(
        'equipment_iot_alerts',
        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True),
        sa.Column('tenant_id', postgresql.UUID(as_uuid=True), nullable=False, index=True),
        sa.Column('equipment_id', postgresql.UUID(as_uuid=True), nullable=False, index=True),
        sa.Column('batch_id', postgresql.UUID(as_uuid=True), nullable=True, index=True),

        # Alert information
        sa.Column('alert_type', sa.String(50), nullable=False),  # temperature_deviation, connection_lost, equipment_error
        sa.Column('severity', sa.String(20), nullable=False),  # info, warning, critical
        sa.Column('alert_time', sa.DateTime(timezone=True), nullable=False, index=True),

        # Alert details
        sa.Column('title', sa.String(255), nullable=False),
        sa.Column('message', sa.Text(), nullable=False),
        sa.Column('sensor_reading_id', postgresql.UUID(as_uuid=True), nullable=True),

        # Threshold information
        sa.Column('threshold_value', sa.Float(), nullable=True),
        sa.Column('actual_value', sa.Float(), nullable=True),
        sa.Column('deviation_percentage', sa.Float(), nullable=True),

        # Status tracking
        sa.Column('is_active', sa.Boolean(), nullable=False, server_default='true'),
        sa.Column('is_acknowledged', sa.Boolean(), nullable=False, server_default='false'),
        sa.Column('acknowledged_by', postgresql.UUID(as_uuid=True), nullable=True),
        sa.Column('acknowledged_at', sa.DateTime(timezone=True), nullable=True),

        sa.Column('is_resolved', sa.Boolean(), nullable=False, server_default='false'),
        sa.Column('resolved_by', postgresql.UUID(as_uuid=True), nullable=True),
        sa.Column('resolved_at', sa.DateTime(timezone=True), nullable=True),
        sa.Column('resolution_notes', sa.Text(), nullable=True),

        # Automated response
        sa.Column('auto_resolved', sa.Boolean(), nullable=False, server_default='false'),
        sa.Column('corrective_action_taken', sa.String(255), nullable=True),

        # Additional data
        sa.Column('additional_data', postgresql.JSON(astext_type=sa.Text()), nullable=True),

        # Timestamps
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.func.now(), onupdate=sa.func.now()),

        # Foreign key constraints
        sa.ForeignKeyConstraint(['equipment_id'], ['equipment.id'], ondelete='CASCADE'),
    )

    # Create indexes for alerts
    op.create_index(
        'idx_iot_alerts_equipment_time',
        'equipment_iot_alerts',
        ['equipment_id', 'alert_time'],
    )
    op.create_index(
        'idx_iot_alerts_active',
        'equipment_iot_alerts',
        ['is_active', 'is_resolved'],
    )


def downgrade():
    """Remove IoT equipment support"""

    # Drop tables
    op.drop_table('equipment_iot_alerts')
    op.drop_table('equipment_connection_logs')
    op.drop_table('equipment_sensor_readings')

    # Remove columns from equipment table
    op.drop_column('equipment', 'supports_remote_control')
    op.drop_column('equipment', 'supports_energy_monitoring')
    op.drop_column('equipment', 'supports_humidity')
    op.drop_column('equipment', 'temperature_zones')
    op.drop_column('equipment', 'poll_interval_seconds')
    op.drop_column('equipment', 'supports_realtime')
    op.drop_column('equipment', 'firmware_version')
    op.drop_column('equipment', 'manufacturer')
    op.drop_column('equipment', 'iot_config')
    op.drop_column('equipment', 'iot_last_connected')
    op.drop_column('equipment', 'iot_connection_status')
    op.drop_column('equipment', 'iot_credentials')
    op.drop_column('equipment', 'iot_port')
    op.drop_column('equipment', 'iot_endpoint')
    op.drop_column('equipment', 'iot_protocol')
    op.drop_column('equipment', 'iot_enabled')
