"""Add missing CASCADE rules and composite indexes for monitoring tables.

Revision ID: 8f9a0b1c2d3e
Revises: 7a8b9c0d1e2f
Create Date: 2026-05-10 19:00:00.000000

"""

from __future__ import annotations

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = "8f9a0b1c2d3e"
down_revision = "7a8b9c0d1e2f"
branch_labels = None
depends_on = None


def upgrade() -> None:
    """Add CASCADE rules and composite indexes."""
    conn = op.get_bind()

    # Use separate transactions for each operation to avoid failure chain
    # First check and drop existing constraints individually
    tables_and_fks = [
        ("cost_snapshots", ["organization_id", "gateway_id"]),
        ("cron_job_statuses", ["organization_id", "gateway_id"]),
        ("session_events", ["organization_id", "gateway_id"]),
        ("sub_agent_runs", ["organization_id", "gateway_id", "session_event_id"]),
        ("system_health_metrics", ["organization_id", "gateway_id"]),
        ("alert_rules", ["organization_id"]),
        ("alert_events", ["organization_id", "rule_id", "gateway_id", "acknowledged_by"]),
    ]

    # Drop existing constraints (ignore if they don't exist)
    for table, columns in tables_and_fks:
        for col in columns:
            constraint_name = f"{table}_{col}_fkey"
            try:
                op.execute(f"ALTER TABLE {table} DROP CONSTRAINT IF EXISTS {constraint_name}")
            except Exception:
                pass

    # Add CASCADE constraints
    # cost_snapshots
    op.execute("ALTER TABLE cost_snapshots ADD CONSTRAINT cost_snapshots_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE cost_snapshots ADD CONSTRAINT cost_snapshots_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id) ON DELETE CASCADE")

    # cron_job_statuses
    op.execute("ALTER TABLE cron_job_statuses ADD CONSTRAINT cron_job_statuses_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE cron_job_statuses ADD CONSTRAINT cron_job_statuses_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id) ON DELETE CASCADE")

    # session_events
    op.execute("ALTER TABLE session_events ADD CONSTRAINT session_events_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE session_events ADD CONSTRAINT session_events_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id) ON DELETE CASCADE")

    # sub_agent_runs
    op.execute("ALTER TABLE sub_agent_runs ADD CONSTRAINT sub_agent_runs_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE sub_agent_runs ADD CONSTRAINT sub_agent_runs_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE sub_agent_runs ADD CONSTRAINT sub_agent_runs_session_event_id_fkey FOREIGN KEY (session_event_id) REFERENCES session_events(id) ON DELETE CASCADE")

    # system_health_metrics
    op.execute("ALTER TABLE system_health_metrics ADD CONSTRAINT system_health_metrics_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE system_health_metrics ADD CONSTRAINT system_health_metrics_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id) ON DELETE CASCADE")

    # alert_rules
    op.execute("ALTER TABLE alert_rules ADD CONSTRAINT alert_rules_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE")

    # alert_events
    op.execute("ALTER TABLE alert_events ADD CONSTRAINT alert_events_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE alert_events ADD CONSTRAINT alert_events_rule_id_fkey FOREIGN KEY (rule_id) REFERENCES alert_rules(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE alert_events ADD CONSTRAINT alert_events_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id) ON DELETE CASCADE")
    op.execute("ALTER TABLE alert_events ADD CONSTRAINT alert_events_acknowledged_by_fkey FOREIGN KEY (acknowledged_by) REFERENCES users(id) ON DELETE SET NULL")

    # Add composite indexes
    op.execute("CREATE INDEX IF NOT EXISTS ix_cost_snapshots_org_created ON cost_snapshots (organization_id, created_at)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_cost_snapshots_org_gateway ON cost_snapshots (organization_id, gateway_id)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_cron_job_statuses_org_created ON cron_job_statuses (organization_id, created_at)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_cron_job_statuses_org_gateway ON cron_job_statuses (organization_id, gateway_id)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_session_events_org_created ON session_events (organization_id, created_at)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_session_events_org_gateway ON session_events (organization_id, gateway_id)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_sub_agent_runs_org_created ON sub_agent_runs (organization_id, created_at)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_sub_agent_runs_org_gateway ON sub_agent_runs (organization_id, gateway_id)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_system_health_metrics_org_created ON system_health_metrics (organization_id, created_at)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_system_health_metrics_org_gateway ON system_health_metrics (organization_id, gateway_id)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_alert_rules_org_created ON alert_rules (organization_id, created_at)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_alert_events_org_created ON alert_events (organization_id, created_at)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_alert_events_org_gateway ON alert_events (organization_id, gateway_id)")
    op.execute("CREATE INDEX IF NOT EXISTS ix_alert_events_rule_created ON alert_events (rule_id, created_at)")


def downgrade() -> None:
    """Remove CASCADE rules and composite indexes."""
    conn = op.get_bind()

    # Drop indexes
    op.execute("DROP INDEX IF EXISTS ix_cost_snapshots_org_created")
    op.execute("DROP INDEX IF EXISTS ix_cost_snapshots_org_gateway")
    op.execute("DROP INDEX IF EXISTS ix_cron_job_statuses_org_created")
    op.execute("DROP INDEX IF EXISTS ix_cron_job_statuses_org_gateway")
    op.execute("DROP INDEX IF EXISTS ix_session_events_org_created")
    op.execute("DROP INDEX IF EXISTS ix_session_events_org_gateway")
    op.execute("DROP INDEX IF EXISTS ix_sub_agent_runs_org_created")
    op.execute("DROP INDEX IF EXISTS ix_sub_agent_runs_org_gateway")
    op.execute("DROP INDEX IF EXISTS ix_system_health_metrics_org_created")
    op.execute("DROP INDEX IF EXISTS ix_system_health_metrics_org_gateway")
    op.execute("DROP INDEX IF EXISTS ix_alert_rules_org_created")
    op.execute("DROP INDEX IF EXISTS ix_alert_events_org_created")
    op.execute("DROP INDEX IF EXISTS ix_alert_events_org_gateway")
    op.execute("DROP INDEX IF EXISTS ix_alert_events_rule_created")

    # Drop constraints (recreate without CASCADE)
    # cost_snapshots
    op.execute("ALTER TABLE cost_snapshots DROP CONSTRAINT IF EXISTS cost_snapshots_organization_id_fkey")
    op.execute("ALTER TABLE cost_snapshots DROP CONSTRAINT IF EXISTS cost_snapshots_gateway_id_fkey")
    op.execute("ALTER TABLE cost_snapshots ADD CONSTRAINT cost_snapshots_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id)")
    op.execute("ALTER TABLE cost_snapshots ADD CONSTRAINT cost_snapshots_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id)")

    # cron_job_statuses
    op.execute("ALTER TABLE cron_job_statuses DROP CONSTRAINT IF EXISTS cron_job_statuses_organization_id_fkey")
    op.execute("ALTER TABLE cron_job_statuses DROP CONSTRAINT IF EXISTS cron_job_statuses_gateway_id_fkey")
    op.execute("ALTER TABLE cron_job_statuses ADD CONSTRAINT cron_job_statuses_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id)")
    op.execute("ALTER TABLE cron_job_statuses ADD CONSTRAINT cron_job_statuses_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id)")

    # session_events
    op.execute("ALTER TABLE session_events DROP CONSTRAINT IF EXISTS session_events_organization_id_fkey")
    op.execute("ALTER TABLE session_events DROP CONSTRAINT IF EXISTS session_events_gateway_id_fkey")
    op.execute("ALTER TABLE session_events ADD CONSTRAINT session_events_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id)")
    op.execute("ALTER TABLE session_events ADD CONSTRAINT session_events_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id)")

    # sub_agent_runs
    op.execute("ALTER TABLE sub_agent_runs DROP CONSTRAINT IF EXISTS sub_agent_runs_organization_id_fkey")
    op.execute("ALTER TABLE sub_agent_runs DROP CONSTRAINT IF EXISTS sub_agent_runs_gateway_id_fkey")
    op.execute("ALTER TABLE sub_agent_runs DROP CONSTRAINT IF EXISTS sub_agent_runs_session_event_id_fkey")
    op.execute("ALTER TABLE sub_agent_runs ADD CONSTRAINT sub_agent_runs_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id)")
    op.execute("ALTER TABLE sub_agent_runs ADD CONSTRAINT sub_agent_runs_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id)")
    op.execute("ALTER TABLE sub_agent_runs ADD CONSTRAINT sub_agent_runs_session_event_id_fkey FOREIGN KEY (session_event_id) REFERENCES session_events(id)")

    # system_health_metrics
    op.execute("ALTER TABLE system_health_metrics DROP CONSTRAINT IF EXISTS system_health_metrics_organization_id_fkey")
    op.execute("ALTER TABLE system_health_metrics DROP CONSTRAINT IF EXISTS system_health_metrics_gateway_id_fkey")
    op.execute("ALTER TABLE system_health_metrics ADD CONSTRAINT system_health_metrics_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id)")
    op.execute("ALTER TABLE system_health_metrics ADD CONSTRAINT system_health_metrics_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id)")

    # alert_rules
    op.execute("ALTER TABLE alert_rules DROP CONSTRAINT IF EXISTS alert_rules_organization_id_fkey")
    op.execute("ALTER TABLE alert_rules ADD CONSTRAINT alert_rules_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id)")

    # alert_events
    op.execute("ALTER TABLE alert_events DROP CONSTRAINT IF EXISTS alert_events_organization_id_fkey")
    op.execute("ALTER TABLE alert_events DROP CONSTRAINT IF EXISTS alert_events_rule_id_fkey")
    op.execute("ALTER TABLE alert_events DROP CONSTRAINT IF EXISTS alert_events_gateway_id_fkey")
    op.execute("ALTER TABLE alert_events DROP CONSTRAINT IF EXISTS alert_events_acknowledged_by_fkey")
    op.execute("ALTER TABLE alert_events ADD CONSTRAINT alert_events_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id)")
    op.execute("ALTER TABLE alert_events ADD CONSTRAINT alert_events_rule_id_fkey FOREIGN KEY (rule_id) REFERENCES alert_rules(id)")
    op.execute("ALTER TABLE alert_events ADD CONSTRAINT alert_events_gateway_id_fkey FOREIGN KEY (gateway_id) REFERENCES gateways(id)")
    op.execute("ALTER TABLE alert_events ADD CONSTRAINT alert_events_acknowledged_by_fkey FOREIGN KEY (acknowledged_by) REFERENCES users(id)")
