"""Add sub_account_id to cmdb_iam_users + extend unique key.

AlicloudIntl is the first provider whose ``list_iam_users`` iterates
sub-accounts (each member of the Resource Directory) and returns RAM
users from each. RAM user names commonly recur across sub-accounts
(``cloud-sso-admin``, ``terraform-runner``, …) — the previous unique
key ``(integration_id, user_name)`` collapsed those into a single row,
silently mis-attributing one sub-account's RAM user as another's. The
mark-and-sweep then nuked the loser, breaking owner-resolution and
offboarding (revoke would target the wrong cloud user).

Fix: add a ``sub_account_id`` column (empty string for non-multi-account
providers / master account) and rebuild the UNIQUE constraint to
include it. NOT NULL with default ``''`` keeps TiDB unique-index
semantics deterministic; AWS / Tencent / single-account paths land
``''`` and behave exactly as before.
"""
from __future__ import annotations

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "120_iam_user_sub_account_unique"
down_revision = "119"
branch_labels = None
depends_on = None


_TABLE = "cmdb_iam_users"
_OLD_UQ = "uq_cmdb_iam_account_user"
_NEW_UQ = "uq_cmdb_iam_account_user_subacct"


def upgrade() -> None:
    bind = op.get_bind()

    # Step 1 — add the column nullable so backfill can run, then make
    # NOT NULL with default '' once populated. Two-phase keeps the
    # online schema change tolerant of TiDB's restrictions on
    # adding NOT-NULL-without-default columns to populated tables.
    op.add_column(
        _TABLE,
        sa.Column("sub_account_id", sa.String(64), nullable=True),
    )

    # Step 2 — backfill from tags JSON ``__sub_account_id`` (set by the
    # AlicloudIntl provider) when present; '' otherwise.
    bind.execute(
        sa.text(
            f"UPDATE {_TABLE} "
            f"SET sub_account_id = COALESCE("
            f"  JSON_UNQUOTE(JSON_EXTRACT(tags, '$.__sub_account_id')), "
            f"  '' "
            f") "
            f"WHERE sub_account_id IS NULL"
        )
    )
    # Treat the JSON-null sentinel ``"null"`` as empty, too.
    bind.execute(
        sa.text(
            f"UPDATE {_TABLE} SET sub_account_id = '' "
            f"WHERE sub_account_id = 'null' OR sub_account_id IS NULL"
        )
    )

    # Step 3 — flip to NOT NULL with default '' so future inserts are
    # safe even if the application code forgets to set it.
    op.alter_column(
        _TABLE,
        "sub_account_id",
        existing_type=sa.String(64),
        nullable=False,
        server_default="",
    )

    # Step 4 — swap the unique constraint. Drop OLD first so the NEW
    # one's index can land without the redundant cover. Tolerate
    # already-applied state (TiDB online schema change can leave the
    # old index in place across retried migrations).
    try:
        op.drop_constraint(_OLD_UQ, _TABLE, type_="unique")
    except Exception:  # noqa: BLE001
        # Some MySQL/TiDB versions report drop_constraint as missing
        # even when present; fall back to the explicit drop_index.
        try:
            op.drop_index(_OLD_UQ, table_name=_TABLE)
        except Exception:  # noqa: BLE001
            pass

    op.create_unique_constraint(
        _NEW_UQ,
        _TABLE,
        ["integration_id", "user_name", "sub_account_id"],
    )


def downgrade() -> None:
    try:
        op.drop_constraint(_NEW_UQ, _TABLE, type_="unique")
    except Exception:  # noqa: BLE001
        try:
            op.drop_index(_NEW_UQ, table_name=_TABLE)
        except Exception:  # noqa: BLE001
            pass
    op.create_unique_constraint(
        _OLD_UQ, _TABLE, ["integration_id", "user_name"],
    )
    op.drop_column(_TABLE, "sub_account_id")