"""add_planning_performance_indexes

Adds performance indexes to optimize plan generation queries.
These indexes significantly reduce query time for capability searches and worker lookups.

Revision ID: b8g9h1i2j3k4
Revises: a7f8e9d1c2b3
Create Date: 2026-01-21 22:00:00.000000

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'b8g9h1i2j3k4'
down_revision: Union[str, Sequence[str], None] = 'df9117888e82'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """
    Add performance indexes for plan generation optimization.

    New indexes:
    1. GIN index on agents.capabilities for fast JSONB searches
    2. Composite index on worker_heartbeats for efficient worker queries

    All indexes are created CONCURRENTLY to avoid locking tables during deployment.
    """

    # Create GIN index on agents.capabilities for JSONB searches
    # This enables fast ILIKE searches on JSONB fields when cast to text
    op.execute("""
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_agents_capabilities_gin
        ON agents USING gin(capabilities jsonb_path_ops)
    """)

    # Create composite index on worker_heartbeats for worker availability queries
    # This optimizes queries that filter by org, queue, and recent heartbeat
    op.execute("""
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_worker_heartbeats_org_queue
        ON worker_heartbeats(organization_id, worker_queue_id, last_heartbeat DESC)
        WHERE status IN ('active', 'idle', 'busy')
    """)


def downgrade() -> None:
    """
    Remove performance indexes.
    """

    # Drop indexes in reverse order
    op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_worker_heartbeats_org_queue")
    op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_agents_capabilities_gin")
