"""expand credit addon validity months to 1-36

Revision ID: b7e2c4d91a63
Revises: 1a4d1f095fa1
Create Date: 2026-08-11 13:15:00
"""

from alembic import op


revision = "b7e2c4d91a63"
down_revision = "1a4d1f095fa1"
branch_labels = None
depends_on = None


CONSTRAINT_NAME = "ck_credit_products_type_required_fields"


def _drop_type_constraint_if_exists() -> None:
    op.execute(
        f"""
        DO $$
        BEGIN
            IF EXISTS (
                SELECT 1
                FROM pg_constraint c
                JOIN pg_class t ON t.oid = c.conrelid
                JOIN pg_namespace n ON n.oid = t.relnamespace
                WHERE n.nspname = current_schema()
                  AND t.relname = 'credit_products'
                  AND c.conname = '{CONSTRAINT_NAME}'
            ) THEN
                ALTER TABLE credit_products DROP CONSTRAINT {CONSTRAINT_NAME};
            END IF;
        END
        $$;
        """
    )


def _create_type_constraint(*, addon_condition: str) -> None:
    op.create_check_constraint(
        CONSTRAINT_NAME,
        "credit_products",
        "(product_type = 'subscription' AND tier_code IS NOT NULL AND tier_rank IS NOT NULL "
        "AND billing_cycle IS NOT NULL AND monthly_grant_credits IS NOT NULL "
        "AND first_purchase_price IS NOT NULL AND regular_price IS NOT NULL "
        "AND grant_credits IS NULL AND validity_months IS NULL) "
        "OR (product_type = 'credit_addon' AND grant_credits IS NOT NULL "
        f"AND {addon_condition} AND tier_code IS NULL AND tier_rank IS NULL "
        "AND billing_cycle IS NULL AND monthly_grant_credits IS NULL "
        "AND first_purchase_price IS NULL AND regular_price IS NULL "
        "AND activity_price IS NULL AND activity_start_at IS NULL AND activity_end_at IS NULL)",
    )


def upgrade() -> None:
    # 当前版本中增值包被旧约束固定为 validity_months = 1。
    # 先移除旧约束，再开放为1-36个自然月；不修改任何现有商品数据。
    _drop_type_constraint_if_exists()
    _create_type_constraint(addon_condition="validity_months BETWEEN 1 AND 36")


def downgrade() -> None:
    # 回退到旧版本时，旧约束只允许1个月。为保证 downgrade 可执行，
    # 将现有增值包有效期恢复为旧版本唯一合法值1个月。
    _drop_type_constraint_if_exists()
    op.execute(
        "UPDATE credit_products "
        "SET validity_months = 1 "
        "WHERE product_type = 'credit_addon' "
        "AND validity_months IS DISTINCT FROM 1"
    )
    _create_type_constraint(addon_condition="validity_months = 1")