"""bailian_usage_daily — per-API-key / per-model chargeback grain.

Revision ID: 299
Revises: 298
Create Date: 2026-06-17

Why this exists
---------------
Bailian (Model Studio) usage used to land as ONE account-level aggregate
row per month, because the ingest read BSS ``QueryBillOverview`` (commodity
level). Verified 2026-06-17: BSS ``DescribeInstanceBill`` (Granularity=
DAILY) returns per-row ``sfm`` cost whose ``InstanceID`` is a composite
``apiKeyId;workspaceId;model;tokenType;channel;0``. The first field is the
global ModelStudio API-key id (matches ``bailian_api_keys.api_key_id``), so
we now persist real per-(day × API key × model) rows and finally do
per-person chargeback. It reconciles to the official bill to the cent.

Schema changes
--------------
1. ``bailian_api_key_id`` (VARCHAR 128, nullable) — the raw ModelStudio
   API-key id parsed from the bill ``InstanceID``. Kept as a string
   (decoupled from the synced ``bailian_api_keys`` FK) so spend on keys
   that live in Resource-Directory MEMBER accounts — which the
   integration AK can't enumerate via ``ListApiKeys`` — is still
   recorded and attributable by account even when ``api_key_pk`` is NULL.
2. ``workspace_id_raw`` (VARCHAR 64, nullable) — the raw workspace id from
   the same composite (member-account workspaces aren't synced into
   ``bailian_workspaces``).
3. Unique key swap: drop ``uq_bailian_usage_integ_date_model_account`` and
   add ``uq_bailian_usage_integ_date_acct_key_model`` over
   ``(integration_id, usage_date, aliyun_account_id, bailian_api_key_id,
   model_name)`` so the upsert keys on the new grain.

Idempotency
-----------
Column / index existence is probed via information_schema, matching the
sibling migrations. Pre-existing aggregate rows (``model_name='aggregate'``,
new columns NULL) don't violate the new unique key (distinct usage_date /
NULL key columns) and are swept by the next ``bailian.usage_sync``.
"""
from __future__ import annotations

import sqlalchemy as sa

from alembic import op

revision = "299"
down_revision = "298"
branch_labels = None
depends_on = None

_TABLE = "bailian_usage_daily"
_OLD_UQ = "uq_bailian_usage_integ_date_model_account"
_NEW_UQ = "uq_bailian_usage_integ_date_acct_key_model"


def _column_exists(bind, table: str, column: str) -> bool:
    return bind.execute(
        sa.text(
            "SELECT 1 FROM information_schema.COLUMNS "
            "WHERE TABLE_SCHEMA = DATABASE() "
            "  AND TABLE_NAME = :t AND COLUMN_NAME = :c LIMIT 1"
        ),
        {"t": table, "c": column},
    ).scalar() is not None


def _index_exists(bind, table: str, index: str) -> bool:
    return bind.execute(
        sa.text(
            "SELECT 1 FROM information_schema.STATISTICS "
            "WHERE TABLE_SCHEMA = DATABASE() "
            "  AND TABLE_NAME = :t AND INDEX_NAME = :i LIMIT 1"
        ),
        {"t": table, "i": index},
    ).scalar() is not None


def upgrade() -> None:
    bind = op.get_bind()
    if not _column_exists(bind, _TABLE, "bailian_api_key_id"):
        op.add_column(
            _TABLE,
            sa.Column("bailian_api_key_id", sa.String(length=128), nullable=True),
        )
    if not _column_exists(bind, _TABLE, "workspace_id_raw"):
        op.add_column(
            _TABLE,
            sa.Column("workspace_id_raw", sa.String(length=64), nullable=True),
        )
    # Swap the unique key onto the new per-key grain. CREATE the new
    # constraint BEFORE dropping the old one: both lead with
    # ``integration_id``, which is the index MySQL/InnoDB uses to support
    # the ``integrations`` foreign key on this table. Dropping the old
    # index first fails with errno 1553 ("needed in a foreign key
    # constraint") because no other index covers ``integration_id`` as a
    # leading column; creating the new one first keeps the FK supported
    # throughout the swap.
    if not _index_exists(bind, _TABLE, _NEW_UQ):
        op.create_unique_constraint(
            _NEW_UQ,
            _TABLE,
            [
                "integration_id",
                "usage_date",
                "aliyun_account_id",
                "bailian_api_key_id",
                "model_name",
            ],
        )
    if _index_exists(bind, _TABLE, _OLD_UQ):
        op.drop_constraint(_OLD_UQ, _TABLE, type_="unique")


def downgrade() -> None:
    bind = op.get_bind()
    # Symmetric to upgrade: recreate the old constraint (leads with
    # integration_id) BEFORE dropping the new one, so the integrations FK
    # stays index-supported throughout (errno 1553 otherwise).
    if not _index_exists(bind, _TABLE, _OLD_UQ):
        op.create_unique_constraint(
            _OLD_UQ,
            _TABLE,
            ["integration_id", "usage_date", "model_name", "aliyun_account_id"],
        )
    if _index_exists(bind, _TABLE, _NEW_UQ):
        op.drop_constraint(_NEW_UQ, _TABLE, type_="unique")
    if _column_exists(bind, _TABLE, "workspace_id_raw"):
        op.drop_column(_TABLE, "workspace_id_raw")
    if _column_exists(bind, _TABLE, "bailian_api_key_id"):
        op.drop_column(_TABLE, "bailian_api_key_id")