"""Add atomic message indexes and storage hot-path indexes.

Revision ID: 0029
Revises: 0028
Create Date: 2026-08-12
"""
from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op


revision: str = "0029"
down_revision: Union[str, None] = "0028"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.add_column(
        "tasks",
        sa.Column("next_message_idx", sa.Integer(), server_default="0", nullable=False),
    )
    op.execute(
        """
        UPDATE tasks AS t
        SET next_message_idx = COALESCE(m.next_idx, 0)
        FROM (
            SELECT task_id, MAX(idx) + 1 AS next_idx
            FROM messages
            GROUP BY task_id
        ) AS m
        WHERE m.task_id = t.task_id
        """
    )

    op.create_index(
        "ix_tasks_user_active_updated",
        "tasks",
        ["user_id", "updated_at", "task_id"],
        postgresql_where=sa.text("deleted_at IS NULL AND scheduled_job_id IS NULL"),
    )
    op.create_index(
        "ix_tasks_schedule_created",
        "tasks",
        ["scheduled_job_id", "created_at", "task_id"],
        postgresql_where=sa.text(
            "deleted_at IS NULL AND scheduled_job_id IS NOT NULL"
        ),
    )
    op.create_index(
        "ix_scheduled_jobs_due_active",
        "scheduled_jobs",
        ["next_run_at"],
        postgresql_where=sa.text("enabled IS TRUE AND deleted_at IS NULL"),
    )
    op.create_index(
        "ix_scheduled_jobs_user_active_created",
        "scheduled_jobs",
        ["user_id", "created_at"],
        postgresql_where=sa.text("deleted_at IS NULL"),
    )
    op.create_index(
        "ix_usage_created_brin",
        "usage_events",
        ["created_at"],
        postgresql_using="brin",
    )


def downgrade() -> None:
    op.drop_index("ix_usage_created_brin", table_name="usage_events")
    op.drop_index(
        "ix_scheduled_jobs_user_active_created", table_name="scheduled_jobs"
    )
    op.drop_index("ix_scheduled_jobs_due_active", table_name="scheduled_jobs")
    op.drop_index("ix_tasks_schedule_created", table_name="tasks")
    op.drop_index("ix_tasks_user_active_updated", table_name="tasks")
    op.drop_column("tasks", "next_message_idx")
