"""add_performance_indexes_v1_0_0

Database optimization migration for mockupAWS v1.0.0
- Composite indexes for frequent queries
- Partial indexes for common filters
- Indexes for N+1 query optimization
- Materialized views for heavy reports

Revision ID: a1b2c3d4e5f6
Revises: efe19595299c
Create Date: 2026-04-07 20:00:00.000000

"""

from typing import Sequence, Union

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


# revision identifiers, used by Alembic.
revision: str = "a1b2c3d4e5f6"
down_revision: Union[str, Sequence[str], None] = "efe19595299c"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema with performance optimizations."""

    # =========================================================================
    # 1. COMPOSITE INDEXES FOR FREQUENT QUERIES
    # =========================================================================

    # Scenario logs: Filter by scenario + date range (common in reports)
    op.create_index(
        "idx_logs_scenario_received",
        "scenario_logs",
        ["scenario_id", "received_at"],
        postgresql_using="btree",
    )

    # Scenario logs: Filter by scenario + source (analytics queries)
    op.create_index(
        "idx_logs_scenario_source",
        "scenario_logs",
        ["scenario_id", "source"],
        postgresql_using="btree",
    )

    # Scenario logs: Filter by scenario + has_pii (PII reports)
    op.create_index(
        "idx_logs_scenario_pii",
        "scenario_logs",
        ["scenario_id", "has_pii"],
        postgresql_using="btree",
    )

    # Scenario logs: Size-based queries (top logs by size)
    op.create_index(
        "idx_logs_scenario_size",
        "scenario_logs",
        ["scenario_id", sa.text("size_bytes DESC")],
        postgresql_using="btree",
    )

    # Scenario metrics: Time-series queries with type filtering
    op.create_index(
        "idx_metrics_scenario_time_type",
        "scenario_metrics",
        ["scenario_id", "timestamp", "metric_type"],
        postgresql_using="btree",
    )

    # Scenario metrics: Name-based aggregation queries
    op.create_index(
        "idx_metrics_scenario_name",
        "scenario_metrics",
        ["scenario_id", "metric_name", "timestamp"],
        postgresql_using="btree",
    )

    # Reports: Scenario + creation date for listing
    op.create_index(
        "idx_reports_scenario_created",
        "reports",
        ["scenario_id", sa.text("created_at DESC")],
        postgresql_using="btree",
    )

    # Scenarios: Status + creation date (dashboard queries)
    op.create_index(
        "idx_scenarios_status_created",
        "scenarios",
        ["status", sa.text("created_at DESC")],
        postgresql_using="btree",
    )

    # Scenarios: Region + status (filtering queries)
    op.create_index(
        "idx_scenarios_region_status",
        "scenarios",
        ["region", "status"],
        postgresql_using="btree",
    )

    # =========================================================================
    # 2. PARTIAL INDEXES FOR COMMON FILTERS
    # =========================================================================

    # Active scenarios only (most queries filter for active)
    op.create_index(
        "idx_scenarios_active",
        "scenarios",
        ["id"],
        postgresql_where=sa.text("status != 'archived'"),
        postgresql_using="btree",
    )

    # Running scenarios (status monitoring)
    op.create_index(
        "idx_scenarios_running",
        "scenarios",
        ["id", "started_at"],
        postgresql_where=sa.text("status = 'running'"),
        postgresql_using="btree",
    )

    # Logs with PII (security audits)
    op.create_index(
        "idx_logs_pii_only",
        "scenario_logs",
        ["scenario_id", "received_at"],
        postgresql_where=sa.text("has_pii = true"),
        postgresql_using="btree",
    )

    # Recent logs index - ordered by received_at DESC for recent queries
    op.create_index(
        "idx_logs_recent",
        "scenario_logs",
        ["scenario_id", sa.text("received_at DESC")],
        postgresql_using="btree",
    )

    # Active API keys
    op.create_index(
        "idx_apikeys_active",
        "api_keys",
        ["user_id", "last_used_at"],
        postgresql_where=sa.text("is_active = true"),
        postgresql_using="btree",
    )

    # Active API keys (valid ones - is_active flag only, can't use NOW() in index predicate)
    op.create_index(
        "idx_apikeys_valid",
        "api_keys",
        ["user_id", "created_at"],
        postgresql_where=sa.text("is_active = true"),
        postgresql_using="btree",
    )

    # =========================================================================
    # 3. INDEXES FOR N+1 QUERY OPTIMIZATION
    # =========================================================================

    # Covering index for scenario list with metrics count
    op.create_index(
        "idx_scenarios_covering",
        "scenarios",
        [
            "id",
            "status",
            "region",
            "created_at",
            "updated_at",
            "total_requests",
            "total_cost_estimate",
        ],
        postgresql_using="btree",
    )

    # Covering index for logs with common fields
    op.create_index(
        "idx_logs_covering",
        "scenario_logs",
        [
            "scenario_id",
            "received_at",
            "source",
            "size_bytes",
            "has_pii",
            "token_count",
        ],
        postgresql_using="btree",
    )

    # =========================================================================
    # 4. ENABLE PG_STAT_STATEMENTS EXTENSION
    # =========================================================================

    op.execute("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")

    # =========================================================================
    # 5. CREATE MATERIALIZED VIEWS FOR HEAVY REPORTS
    # =========================================================================

    # Daily scenario statistics (refreshed nightly)
    op.execute("""
        CREATE MATERIALIZED VIEW IF NOT EXISTS mv_scenario_daily_stats AS
        SELECT 
            s.id as scenario_id,
            s.name as scenario_name,
            s.status,
            s.region,
            DATE(sl.received_at) as log_date,
            COUNT(sl.id) as log_count,
            SUM(sl.size_bytes) as total_size_bytes,
            SUM(sl.token_count) as total_tokens,
            SUM(sl.sqs_blocks) as total_sqs_blocks,
            COUNT(CASE WHEN sl.has_pii THEN 1 END) as pii_count,
            COUNT(DISTINCT sl.source) as unique_sources
        FROM scenarios s
        LEFT JOIN scenario_logs sl ON s.id = sl.scenario_id
        WHERE sl.received_at > NOW() - INTERVAL '90 days'
        GROUP BY s.id, s.name, s.status, s.region, DATE(sl.received_at)
        ORDER BY log_date DESC
    """)

    op.create_index(
        "idx_mv_daily_stats_scenario",
        "mv_scenario_daily_stats",
        ["scenario_id", "log_date"],
        postgresql_using="btree",
    )

    # Monthly cost aggregation
    op.execute("""
        CREATE MATERIALIZED VIEW IF NOT EXISTS mv_monthly_costs AS
        SELECT 
            DATE_TRUNC('month', sm.timestamp) as month,
            sm.scenario_id,
            sm.metric_type,
            sm.metric_name,
            SUM(sm.value) as total_value,
            AVG(sm.value)::numeric(15,6) as avg_value,
            MAX(sm.value)::numeric(15,6) as max_value,
            MIN(sm.value)::numeric(15,6) as min_value,
            COUNT(*) as metric_count
        FROM scenario_metrics sm
        WHERE sm.timestamp > NOW() - INTERVAL '2 years'
        GROUP BY DATE_TRUNC('month', sm.timestamp), sm.scenario_id, sm.metric_type, sm.metric_name
        ORDER BY month DESC
    """)

    op.create_index(
        "idx_mv_monthly_costs_lookup",
        "mv_monthly_costs",
        ["scenario_id", "month", "metric_type"],
        postgresql_using="btree",
    )

    # Source analytics summary
    op.execute("""
        CREATE MATERIALIZED VIEW IF NOT EXISTS mv_source_analytics AS
        SELECT 
            sl.scenario_id,
            sl.source,
            DATE_TRUNC('day', sl.received_at) as day,
            COUNT(*) as log_count,
            SUM(sl.size_bytes) as total_bytes,
            AVG(sl.size_bytes)::numeric(12,2) as avg_size_bytes,
            SUM(sl.token_count) as total_tokens,
            AVG(sl.token_count)::numeric(12,2) as avg_tokens,
            COUNT(CASE WHEN sl.has_pii THEN 1 END) as pii_count
        FROM scenario_logs sl
        WHERE sl.received_at > NOW() - INTERVAL '30 days'
        GROUP BY sl.scenario_id, sl.source, DATE_TRUNC('day', sl.received_at)
        ORDER BY day DESC, log_count DESC
    """)

    op.create_index(
        "idx_mv_source_analytics_lookup",
        "mv_source_analytics",
        ["scenario_id", "day"],
        postgresql_using="btree",
    )

    # =========================================================================
    # 6. CREATE REFRESH FUNCTION FOR MATERIALIZED VIEWS
    # =========================================================================

    op.execute("""
        CREATE OR REPLACE FUNCTION refresh_materialized_views()
        RETURNS void AS $$
        BEGIN
            REFRESH MATERIALIZED VIEW CONCURRENTLY mv_scenario_daily_stats;
            REFRESH MATERIALIZED VIEW CONCURRENTLY mv_monthly_costs;
            REFRESH MATERIALIZED VIEW CONCURRENTLY mv_source_analytics;
        END;
        $$ LANGUAGE plpgsql
    """)

    # =========================================================================
    # 7. CREATE QUERY PERFORMANCE LOGGING TABLE
    # =========================================================================

    op.create_table(
        "query_performance_log",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            primary_key=True,
            server_default=sa.text("uuid_generate_v4()"),
        ),
        sa.Column("query_hash", sa.String(64), nullable=False),
        sa.Column("query_text", sa.Text(), nullable=False),
        sa.Column("execution_time_ms", sa.Integer(), nullable=False),
        sa.Column("rows_affected", sa.Integer(), nullable=True),
        sa.Column(
            "created_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("NOW()"),
            nullable=False,
        ),
        sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=True),
        sa.Column("endpoint", sa.String(255), nullable=True),
    )

    op.create_index(
        "idx_query_perf_hash",
        "query_performance_log",
        ["query_hash"],
        postgresql_using="btree",
    )

    op.create_index(
        "idx_query_perf_time",
        "query_performance_log",
        ["created_at"],
        postgresql_using="brin",
    )

    op.create_index(
        "idx_query_perf_slow",
        "query_performance_log",
        ["execution_time_ms"],
        postgresql_where=sa.text("execution_time_ms > 1000"),
        postgresql_using="btree",
    )


def downgrade() -> None:
    """Downgrade schema."""

    # Drop query performance log table
    op.drop_index("idx_query_perf_slow", table_name="query_performance_log")
    op.drop_index("idx_query_perf_time", table_name="query_performance_log")
    op.drop_index("idx_query_perf_hash", table_name="query_performance_log")
    op.drop_table("query_performance_log")

    # Drop refresh function
    op.execute("DROP FUNCTION IF EXISTS refresh_materialized_views()")

    # Drop materialized views
    op.drop_index("idx_mv_source_analytics_lookup", table_name="mv_source_analytics")
    op.execute("DROP MATERIALIZED VIEW IF EXISTS mv_source_analytics")

    op.drop_index("idx_mv_monthly_costs_lookup", table_name="mv_monthly_costs")
    op.execute("DROP MATERIALIZED VIEW IF EXISTS mv_monthly_costs")

    op.drop_index("idx_mv_daily_stats_scenario", table_name="mv_scenario_daily_stats")
    op.execute("DROP MATERIALIZED VIEW IF EXISTS mv_scenario_daily_stats")

    # Drop indexes (composite)
    op.drop_index("idx_scenarios_region_status", table_name="scenarios")
    op.drop_index("idx_scenarios_status_created", table_name="scenarios")
    op.drop_index("idx_reports_scenario_created", table_name="reports")
    op.drop_index("idx_metrics_scenario_name", table_name="scenario_metrics")
    op.drop_index("idx_metrics_scenario_time_type", table_name="scenario_metrics")
    op.drop_index("idx_logs_scenario_size", table_name="scenario_logs")
    op.drop_index("idx_logs_scenario_pii", table_name="scenario_logs")
    op.drop_index("idx_logs_scenario_source", table_name="scenario_logs")
    op.drop_index("idx_logs_scenario_received", table_name="scenario_logs")

    # Drop indexes (partial)
    op.drop_index("idx_apikeys_valid", table_name="api_keys")
    op.drop_index("idx_apikeys_active", table_name="api_keys")
    op.drop_index("idx_logs_recent", table_name="scenario_logs")
    op.drop_index("idx_logs_pii_only", table_name="scenario_logs")
    op.drop_index("idx_scenarios_running", table_name="scenarios")
    op.drop_index("idx_scenarios_active", table_name="scenarios")

    # Drop indexes (covering)
    op.drop_index("idx_logs_covering", table_name="scenario_logs")
    op.drop_index("idx_scenarios_covering", table_name="scenarios")
