"""drop credit ratio model config foreign key and add engine indexes

Revision ID: 6846d43389b6
Revises: c310d7193c7f
Create Date: 2026-05-28 11:04:37.270041
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = "6846d43389b6"
down_revision: Union[str, None] = "c310d7193c7f"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


OLD_FK_NAME = "credit_ratios_model_config_id_fkey"


def _index_exists(bind, table_name: str, index_name: str) -> bool:
    inspector = sa.inspect(bind)
    return any(index.get("name") == index_name for index in inspector.get_indexes(table_name))


def _drop_index_if_exists(bind, index_name: str, table_name: str) -> None:
    if _index_exists(bind, table_name, index_name):
        op.drop_index(index_name, table_name=table_name)


def _create_index_if_missing(
    bind,
    index_name: str,
    table_name: str,
    columns: list[str],
    unique: bool = False,
) -> None:
    if not _index_exists(bind, table_name, index_name):
        op.create_index(index_name, table_name, columns, unique=unique)


def _get_credit_ratio_model_config_fk_name(bind) -> str | None:
    """兼容不同数据库/命名规则，查找 credit_ratios.model_config_id -> model_configs.id 的外键名。"""
    inspector = sa.inspect(bind)
    for fk in inspector.get_foreign_keys("credit_ratios"):
        constrained_columns = fk.get("constrained_columns") or []
        referred_table = fk.get("referred_table")
        referred_columns = fk.get("referred_columns") or []
        fk_name = fk.get("name")

        if (
            constrained_columns == ["model_config_id"]
            and referred_table == "model_configs"
            and referred_columns == ["id"]
        ):
            return fk_name

    return None


def _table_has_rows(bind, table_name: str, where_sql: str | None = None) -> bool:
    sql = f"SELECT COUNT(*) FROM {table_name}"
    if where_sql:
        sql += f" WHERE {where_sql}"
    count = bind.execute(sa.text(sql)).scalar()
    return bool(count)


def _get_default_video_engine_id(bind) -> str | None:
    """获取默认视频引擎ID：优先启用状态 priority 降序第一；没有启用时取全部 priority 降序第一。"""
    engine_id = bind.execute(
        sa.text(
            """
            SELECT id
            FROM video_engines
            WHERE is_active IS TRUE
            ORDER BY priority DESC, id ASC
            LIMIT 1
            """
        )
    ).scalar()

    if engine_id:
        return engine_id

    return bind.execute(
        sa.text(
            """
            SELECT id
            FROM video_engines
            ORDER BY priority DESC, id ASC
            LIMIT 1
            """
        )
    ).scalar()


def _get_default_image_engine_id(bind) -> str | None:
    """获取默认图片引擎ID：优先启用状态 priority 降序第一；没有启用时取全部 priority 降序第一。"""
    engine_id = bind.execute(
        sa.text(
            """
            SELECT id
            FROM image_engines
            WHERE is_active IS TRUE
            ORDER BY priority DESC, id ASC
            LIMIT 1
            """
        )
    ).scalar()

    if engine_id:
        return engine_id

    return bind.execute(
        sa.text(
            """
            SELECT id
            FROM image_engines
            ORDER BY priority DESC, id ASC
            LIMIT 1
            """
        )
    ).scalar()


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

    # 1. 删除 credit_ratios.model_config_id -> model_configs.id 外键。
    # 删除后 model_config_id 保留旧字段名，但业务含义变更为：
    # - gen_type=video 时保存 video_engines.id
    # - gen_type=image 时保存 image_engines.id
    fk_name = _get_credit_ratio_model_config_fk_name(bind)
    if fk_name:
        op.drop_constraint(fk_name, "credit_ratios", type_="foreignkey")

    # 2. 规范历史 gen_type，避免 Video / IMAGE / 空格 影响后续数据更新。
    bind.execute(
        sa.text(
            """
            UPDATE credit_ratios
            SET gen_type = LOWER(TRIM(gen_type))
            WHERE gen_type IS NOT NULL
            """
        )
    )

    # 3. 将当前 CreditRatio 内原来绑定 ModelConfig 的 model_config_id
    #    改为不同类型对应的默认引擎ID。
    #
    #    gen_type=video -> video_engines 中 priority 权重降序第一的 ID
    #    gen_type=image -> image_engines 中 priority 权重降序第一的 ID
    #
    #    这里不建备份表，降级也不回滚这部分业务数据。
    default_video_engine_id = _get_default_video_engine_id(bind)
    default_image_engine_id = _get_default_image_engine_id(bind)

    has_video_ratios = _table_has_rows(bind, "credit_ratios", "gen_type = 'video'")
    has_image_ratios = _table_has_rows(bind, "credit_ratios", "gen_type = 'image'")

    if has_video_ratios and not default_video_engine_id:
        raise RuntimeError(
            "迁移失败：credit_ratios 中存在 gen_type=video 的积分规则，"
            "但 video_engines 表没有可用视频引擎，无法将 model_config_id 改为默认视频引擎ID。"
        )

    if has_image_ratios and not default_image_engine_id:
        raise RuntimeError(
            "迁移失败：credit_ratios 中存在 gen_type=image 的积分规则，"
            "但 image_engines 表没有可用图片引擎，无法将 model_config_id 改为默认图片引擎ID。"
        )

    if default_video_engine_id:
        bind.execute(
            sa.text(
                """
                UPDATE credit_ratios
                SET model_config_id = :engine_id
                WHERE gen_type = 'video'
                """
            ),
            {"engine_id": default_video_engine_id},
        )

    if default_image_engine_id:
        bind.execute(
            sa.text(
                """
                UPDATE credit_ratios
                SET model_config_id = :engine_id
                WHERE gen_type = 'image'
                """
            ),
            {"engine_id": default_image_engine_id},
        )

    # 4. 创建当前模型层需要的索引。
    _create_index_if_missing(
        bind,
        "ix_credit_ratios_gen_type",
        "credit_ratios",
        ["gen_type"],
    )
    _create_index_if_missing(
        bind,
        "ix_credit_ratios_model_config_id",
        "credit_ratios",
        ["model_config_id"],
    )
    _create_index_if_missing(
        bind,
        "ix_credit_ratios_resolution",
        "credit_ratios",
        ["resolution"],
    )
    _create_index_if_missing(
        bind,
        "ix_credit_ratios_gen_type_engine_resolution",
        "credit_ratios",
        ["gen_type", "model_config_id", "resolution"],
    )
    _create_index_if_missing(
        bind,
        "ix_credit_ratios_gen_type_resolution",
        "credit_ratios",
        ["gen_type", "resolution"],
    )


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

    # 业务需求：降级不回滚 credit_ratios.model_config_id 数据。
    # 当前 model_config_id 已经是 image_engines.id / video_engines.id，
    # 因此这里也不恢复 credit_ratios.model_config_id -> model_configs.id 外键。
    #
    # 否则会出现两种问题：
    # 1. 字段里是引擎ID，不是 model_configs.id，恢复外键会失败；
    # 2. 如果为了恢复外键强行回填 model_configs.id，就违背“降级不回滚业务数据”的要求。
    #
    # 所以 downgrade 只撤销本次新增索引，保留业务数据和取消外键后的结构。
    _drop_index_if_exists(bind, "ix_credit_ratios_gen_type_resolution", "credit_ratios")
    _drop_index_if_exists(bind, "ix_credit_ratios_gen_type_engine_resolution", "credit_ratios")
    _drop_index_if_exists(bind, "ix_credit_ratios_resolution", "credit_ratios")
    _drop_index_if_exists(bind, "ix_credit_ratios_model_config_id", "credit_ratios")
    _drop_index_if_exists(bind, "ix_credit_ratios_gen_type", "credit_ratios")