"""expand commercial runtime resource quantity bases

Revision ID: 20260717_0024
Revises: 20260716_0023
Create Date: 2026-07-17 09:30:00
"""

from collections.abc import Sequence

import sqlalchemy as sa

from alembic import op

revision: str = "20260717_0024"
down_revision: str | None = "20260716_0023"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

_CONSTRAINT = "ck_commercial_runtime_reservations_basis"
_LEGACY_BASES = (
    "quantity_basis IN ('call', 'input_tokens', 'output_tokens', "
    "'total_tokens', 'duration_ms')"
)
_RESOURCE_BASES = (
    "quantity_basis IN ('call', 'input_tokens', 'output_tokens', "
    "'total_tokens', 'duration_ms', 'bytes', 'pages', 'objects', 'events')"
)


def _require_postgresql() -> None:
    dialect_name = op.get_bind().dialect.name
    if dialect_name != "postgresql":
        raise RuntimeError(
            "20260717_0024 only supports PostgreSQL; "
            f"refusing to mutate {dialect_name} without transactional constraint DDL"
        )


def _require_no_resource_reservations_for_downgrade() -> None:
    count = int(
        op.get_bind().scalar(
            sa.text(
                "SELECT COUNT(*) FROM commercial_runtime_reservations "
                "WHERE quantity_basis IN ('bytes', 'pages', 'objects', 'events')"
            )
        )
        or 0
    )
    if count:
        raise RuntimeError(
            "cannot downgrade commercial resource quantity bases: "
            f"resource reservations exist ({count})"
        )


def upgrade() -> None:
    _require_postgresql()
    op.drop_constraint(
        _CONSTRAINT,
        "commercial_runtime_reservations",
        type_="check",
    )
    op.create_check_constraint(
        _CONSTRAINT,
        "commercial_runtime_reservations",
        _RESOURCE_BASES,
    )


def downgrade() -> None:
    _require_postgresql()
    _require_no_resource_reservations_for_downgrade()
    op.drop_constraint(
        _CONSTRAINT,
        "commercial_runtime_reservations",
        type_="check",
    )
    op.create_check_constraint(
        _CONSTRAINT,
        "commercial_runtime_reservations",
        _LEGACY_BASES,
    )