"""add credit_record biz_key

Revision ID: 516b84e3bf17
Revises: e7bf423b248c
Create Date: 2026-06-03 16:12:46.784240
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


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


def upgrade() -> None:
    # 1. CreditRecord 新增正式幂等字段
    op.add_column(
        "credit_records",
        sa.Column("biz_key", sa.String(length=160), nullable=True),
    )
    op.add_column(
        "credit_records",
        sa.Column("refund_for_biz_key", sa.String(length=160), nullable=True),
    )

    # 2. 普通查询索引
    op.create_index(
        "ix_credit_records_biz_key",
        "credit_records",
        ["biz_key"],
        unique=False,
    )
    op.create_index(
        "ix_credit_records_refund_for_biz_key",
        "credit_records",
        ["refund_for_biz_key"],
        unique=False,
    )
    op.create_index(
        "ix_credit_records_related_type",
        "credit_records",
        ["related_id", "type"],
        unique=False,
    )
    op.create_index(
        "ix_credit_records_user_refund_for_biz_key",
        "credit_records",
        ["user_id", "refund_for_biz_key"],
        unique=False,
    )

    # 3. PostgreSQL 部分唯一索引：只限制 biz_key 不为空的记录
    op.create_index(
        "uq_credit_records_user_biz_key",
        "credit_records",
        ["user_id", "biz_key"],
        unique=True,
        postgresql_where=sa.text("biz_key IS NOT NULL"),
    )

    # 4. ChatGenerationTask 幂等 key 唯一索引
    # 只限制 idempotency_key 不为空且不为空字符串的数据
    op.create_index(
        "uq_chat_generation_tasks_user_mode_idempotency",
        "chat_generation_tasks",
        ["user_id", "generation_mode", "idempotency_key"],
        unique=True,
        postgresql_where=sa.text(
            "idempotency_key IS NOT NULL AND idempotency_key <> ''"
        ),
    )


def downgrade() -> None:
    op.drop_index(
        "uq_chat_generation_tasks_user_mode_idempotency",
        table_name="chat_generation_tasks",
    )
    op.drop_index(
        "uq_credit_records_user_biz_key",
        table_name="credit_records",
    )
    op.drop_index(
        "ix_credit_records_user_refund_for_biz_key",
        table_name="credit_records",
    )
    op.drop_index(
        "ix_credit_records_related_type",
        table_name="credit_records",
    )
    op.drop_index(
        "ix_credit_records_refund_for_biz_key",
        table_name="credit_records",
    )
    op.drop_index(
        "ix_credit_records_biz_key",
        table_name="credit_records",
    )
    op.drop_column("credit_records", "refund_for_biz_key")
    op.drop_column("credit_records", "biz_key")