Aether/_deprecated_py_src/services/model/global_model.py

"""
GlobalModel 服务层

提供 GlobalModel 的 CRUD 操作、查询和统计功能
"""

from __future__ import annotations

from typing import cast

from sqlalchemy import delete as sa_delete
from sqlalchemy import func
from sqlalchemy.orm import Session, joinedload, load_only

from src.core.exceptions import InvalidRequestException, NotFoundException
from src.core.logger import logger
from src.models.database import GlobalModel, Model
from src.models.pydantic_models import GlobalModelUpdate


async def on_key_allowed_models_changed(
    db: Session,
    provider_id: str,
    allowed_models: list[str] | None = None,
    skip_disassociate: bool = False,
) -> None:
    """
    Key 的 allowed_models 变更后的统一处理

    包括：
    1. 触发缓存失效（包括 /v1/models 列表缓存）
    2. 检查并自动关联匹配的 GlobalModel（仅当提供 allowed_models 时）
    3. 检查并自动解除不再匹配的 GlobalModel 关联（可通过 skip_disassociate 跳过）

    Args:
        db: 数据库 Session
        provider_id: Provider ID
        allowed_models: 更新后的 allowed_models 列表
            - 提供非空列表：触发自动关联和解除关联检查
            - 提供空列表或 None：仅触发解除关联检查（用于 Key 删除场景）
        skip_disassociate: 是否跳过解除关联检查
            - True：跳过（用于删除 allowed_models 为 null 的 Key 时）
            - False：执行检查（默认）
    """
    from src.services.cache.invalidation import get_cache_invalidation_service

    # 1. 触发缓存失效
    cache_service = get_cache_invalidation_service()
    await cache_service.on_key_allowed_models_changed(provider_id)

    # 2. 检查并自动关联 GlobalModel（仅当提供非空 allowed_models 时）
    if allowed_models:
        GlobalModelService.auto_associate_provider_by_key_whitelist(
            db=db,
            provider_id=provider_id,
            allowed_models=allowed_models,
        )

    # 3. 检查并自动解除不再匹配的 GlobalModel 关联
    if not skip_disassociate:
        GlobalModelService.auto_disassociate_provider_by_key_whitelist(
            db=db,
            provider_id=provider_id,
        )


class GlobalModelService:
    """GlobalModel 服务"""

    @staticmethod
    def get_global_model(db: Session, global_model_id: str) -> GlobalModel:
        """
        获取单个 GlobalModel

        Args:
            global_model_id: GlobalModel 的 UUID 或 name
        """
        # 先尝试通过 ID 查找
        global_model = db.query(GlobalModel).filter(GlobalModel.id == global_model_id).first()

        # 如果没找到，尝试通过 name 查找
        if not global_model:
            global_model = db.query(GlobalModel).filter(GlobalModel.name == global_model_id).first()

        if not global_model:
            raise NotFoundException(f"GlobalModel {global_model_id} not found")
        return global_model

    @staticmethod
    def get_global_model_by_name(db: Session, name: str) -> GlobalModel | None:
        """通过名称获取 GlobalModel"""
        return db.query(GlobalModel).filter(GlobalModel.name == name).first()

    @staticmethod
    def list_global_models(
        db: Session,
        skip: int = 0,
        limit: int = 100,
        is_active: bool | None = None,
        search: str | None = None,
    ) -> list[GlobalModel]:
        """列出 GlobalModel"""
        query = db.query(GlobalModel)

        if is_active is not None:
            query = query.filter(GlobalModel.is_active == is_active)

        if search:
            search_pattern = f"%{search}%"
            query = query.filter(
                (GlobalModel.name.ilike(search_pattern))
                | (GlobalModel.display_name.ilike(search_pattern))
            )

        # 按名称排序
        query = query.order_by(GlobalModel.name)

        return query.offset(skip).limit(limit).all()

    @staticmethod
    def create_global_model(
        db: Session,
        name: str,
        display_name: str,
        is_active: bool | None = True,
        # 按次计费配置
        default_price_per_request: float | None = None,
        # 阶梯计费配置（必填）
        default_tiered_pricing: dict | None = None,
        # Key 能力配置
        supported_capabilities: list[str] | None = None,
        # 模型配置（JSON）
        config: dict | None = None,
    ) -> GlobalModel:
        """创建 GlobalModel"""
        # 检查名称是否已存在
        existing = GlobalModelService.get_global_model_by_name(db, name)
        if existing:
            raise InvalidRequestException(f"GlobalModel with name '{name}' already exists")

        global_model = GlobalModel(
            name=name,
            display_name=display_name,
            is_active=is_active,
            # 按次计费配置
            default_price_per_request=default_price_per_request,
            # 阶梯计费配置
            default_tiered_pricing=default_tiered_pricing,
            # Key 能力配置
            supported_capabilities=supported_capabilities,
            # 模型配置（JSON）
            config=config,
        )

        db.add(global_model)
        db.commit()
        db.refresh(global_model)

        return global_model

    @staticmethod
    def update_global_model(
        db: Session,
        global_model_id: str,
        update_data: GlobalModelUpdate,
    ) -> GlobalModel:
        """
        更新 GlobalModel

        使用 exclude_unset=True 来区分"未提供字段"和"显式设置为 None"：
        - 未提供的字段不会被更新
        - 显式设置为 None 的字段会被更新为 None（置空）
        """
        global_model = GlobalModelService.get_global_model(db, global_model_id)

        # 只更新显式设置的字段（包括显式设置为 None 的情况）
        data_dict = update_data.model_dump(exclude_unset=True)

        # 处理阶梯计费配置：如果是 TieredPricingConfig 对象，转换为 dict
        if "default_tiered_pricing" in data_dict:
            tiered_pricing = data_dict["default_tiered_pricing"]
            if tiered_pricing is not None and hasattr(tiered_pricing, "model_dump"):
                data_dict["default_tiered_pricing"] = tiered_pricing.model_dump()

        for field, value in data_dict.items():
            setattr(global_model, field, value)

        db.commit()
        db.refresh(global_model)

        return global_model

    @staticmethod
    def delete_global_model(db: Session, global_model_id: str) -> None:
        """
        删除 GlobalModel

        默认行为: 级联删除所有关联的 Provider 模型实现
        注意: 不清理 API Key 和 User 的 allowed_models 引用，
        保留无效引用可让用户在前端看到"已失效"的模型，便于手动清理或等待重建同名模型
        """
        global_model = GlobalModelService.get_global_model(db, global_model_id)

        # 批量删除所有关联的 Provider 模型实现
        assoc_count = (
            db.query(func.count(Model.id)).filter(Model.global_model_id == global_model.id).scalar()
        )
        if assoc_count:
            logger.info(
                f"删除 GlobalModel {global_model.name} 的 {assoc_count} 个关联 Provider 模型"
            )
            db.execute(sa_delete(Model).where(Model.global_model_id == global_model.id))

        # 删除 GlobalModel
        db.delete(global_model)
        db.commit()

    @staticmethod
    def get_global_model_stats(db: Session, global_model_id: str) -> dict:
        """获取 GlobalModel 统计信息"""
        global_model = GlobalModelService.get_global_model(db, global_model_id)

        # 统计关联的 Model 数量（使用 global_model.id，预加载 provider 关联）
        models = (
            db.query(Model)
            .options(joinedload(Model.provider))
            .filter(Model.global_model_id == global_model.id)
            .all()
        )

        # 统计支持的 Provider 数量
        provider_ids = {model.provider_id for model in models}

        # 从阶梯计费中提取价格范围
        input_prices = []
        output_prices = []
        for m in models:
            tiered = m.get_effective_tiered_pricing()
            if tiered and tiered.get("tiers"):
                first_tier = tiered["tiers"][0]
                if first_tier.get("input_price_per_1m") is not None:
                    input_prices.append(first_tier["input_price_per_1m"])
                if first_tier.get("output_price_per_1m") is not None:
                    output_prices.append(first_tier["output_price_per_1m"])

        return {
            "global_model_id": global_model.id,
            "name": global_model.name,
            "total_models": len(models),
            "total_providers": len(provider_ids),
            "price_range": {
                "min_input": min(input_prices) if input_prices else None,
                "max_input": max(input_prices) if input_prices else None,
                "min_output": min(output_prices) if output_prices else None,
                "max_output": max(output_prices) if output_prices else None,
            },
        }

    @staticmethod
    def batch_assign_to_providers(
        db: Session,
        global_model_id: str,
        provider_ids: list[str],
        create_models: bool = False,
    ) -> dict:
        """批量为多个 Provider 添加 GlobalModel 实现"""

        global_model = GlobalModelService.get_global_model(db, global_model_id)

        results = {
            "success": [],
            "errors": [],
        }

        for provider_id in provider_ids:
            try:
                # 检查该 Provider 是否已有该 GlobalModel 的实现（使用 global_model.id）
                existing_model = (
                    db.query(Model)
                    .filter(
                        Model.provider_id == provider_id,
                        Model.global_model_id == global_model.id,
                    )
                    .first()
                )

                if existing_model:
                    results["errors"].append(
                        {
                            "provider_id": provider_id,
                            "error": "Model already exists for this provider",
                        }
                    )
                    continue

                if create_models:
                    # 创建新的 Model（价格和能力设为 None，继承 GlobalModel 默认值）
                    model = Model(
                        provider_id=provider_id,
                        global_model_id=global_model.id,
                        provider_model_name=global_model.name,  # 默认使用 GlobalModel name
                        # 计费设为 None，使用 GlobalModel 默认值
                        price_per_request=None,
                        tiered_pricing=None,
                        # 能力设为 None，使用 GlobalModel 默认值
                        supports_vision=None,
                        supports_function_calling=None,
                        supports_streaming=None,
                        supports_extended_thinking=None,
                        is_active=True,
                    )
                    db.add(model)
                    db.commit()

                    results["success"].append(
                        {"provider_id": provider_id, "model_id": model.id, "created": True}
                    )
                else:
                    results["errors"].append(
                        {
                            "provider_id": provider_id,
                            "error": "create_models=False, no existing model found",
                        }
                    )

            except Exception as e:
                db.rollback()
                results["errors"].append({"provider_id": provider_id, "error": str(e)})

        db.commit()
        return results

    @staticmethod
    def auto_associate_provider_by_key_whitelist(
        db: Session,
        provider_id: str,
        allowed_models: list[str],
    ) -> dict:
        """
        根据 Key 白名单自动关联 Provider 到匹配的 GlobalModel

        当 Key 的 allowed_models 更新后调用此方法，检查所有 GlobalModel 的映射规则，
        如果有映射规则匹配到 Key 白名单中的模型，且 Provider 尚未关联到该 GlobalModel，
        则自动创建关联。

        Args:
            db: 数据库 Session
            provider_id: Provider ID
            allowed_models: Key 的白名单模型列表

        Returns:
            Dict: 包含 success 和 errors 列表
        """
        from src.core.model_permissions import match_model_with_pattern
        from src.models.database import Provider

        results: dict[str, list[dict]] = {
            "success": [],
            "errors": [],
        }

        if not allowed_models:
            return results

        # 获取 Provider
        provider = db.query(Provider).filter(Provider.id == provider_id).first()
        if not provider:
            logger.warning(f"Provider {provider_id} not found for auto-association")
            return results

        # 获取该 Provider 已关联的 GlobalModel ID 集合
        existing_associations = (
            db.query(Model.global_model_id, Model.provider_model_name)
            .filter(Model.provider_id == provider_id)
            .all()
        )
        linked_global_model_ids: set[str] = {row[0] for row in existing_associations if row[0]}
        # 同时获取已存在的 provider_model_name 集合，避免唯一约束冲突
        existing_provider_model_names: set[str] = {
            row[1] for row in existing_associations if row[1]
        }

        # 获取所有活跃的 GlobalModel（带映射规则）
        global_models = db.query(GlobalModel).filter(GlobalModel.is_active == True).all()

        allowed_models_set = set(allowed_models)

        for global_model in global_models:
            # 跳过已关联的
            if global_model.id in linked_global_model_ids:
                continue

            # 跳过 provider_model_name 已存在的（避免唯一约束冲突）
            if global_model.name in existing_provider_model_names:
                logger.debug(
                    f"Skipping auto-association for GlobalModel {global_model.name}: "
                    f"provider_model_name already exists for Provider {provider.name}"
                )
                continue

            # 提取映射规则
            model_mappings: list[str] = []
            if global_model.config and isinstance(global_model.config, dict):
                mappings = global_model.config.get("model_mappings")
                if isinstance(mappings, list):
                    model_mappings = [m for m in mappings if isinstance(m, str)]

            if not model_mappings:
                continue

            # 检查是否有映射规则匹配到 Key 白名单
            matched = False
            for mapping_pattern in model_mappings:
                for allowed_model in allowed_models_set:
                    if match_model_with_pattern(mapping_pattern, allowed_model):
                        matched = True
                        break
                if matched:
                    break

            if not matched:
                continue

            # 自动创建关联（逐个处理，允许部分成功）
            try:
                new_model = Model(
                    provider_id=provider_id,
                    global_model_id=global_model.id,
                    provider_model_name=global_model.name,
                    is_active=True,
                )
                db.add(new_model)
                db.flush()

                # 添加到已存在集合，避免后续循环重复创建
                existing_provider_model_names.add(global_model.name)

                results["success"].append(
                    {
                        "global_model_id": global_model.id,
                        "global_model_name": global_model.name,
                        "model_id": new_model.id,
                    }
                )
                logger.info(
                    f"Auto-associated Provider {provider.name} to GlobalModel {global_model.name} "
                    f"via mapping rule match"
                )
            except Exception as e:
                db.rollback()
                logger.error(
                    f"Failed to auto-associate Provider {provider.name} to GlobalModel {global_model.name}: {e}"
                )
                results["errors"].append(
                    {
                        "global_model_id": global_model.id,
                        "global_model_name": global_model.name,
                        "error": str(e),
                    }
                )

        if results["success"]:
            db.commit()

        return results

    @staticmethod
    def auto_disassociate_provider_by_key_whitelist(
        db: Session,
        provider_id: str,
    ) -> dict:
        """
        根据 Key 白名单自动解除 Provider 与不再匹配的 GlobalModel 的关联

        当 Key 的 allowed_models 更新后调用此方法，检查所有已关联的 GlobalModel，
        如果其映射规则不再匹配任何 Key 白名单中的模型，则自动删除关联。

        注意：只删除通过映射规则自动关联的 Model（即 GlobalModel 有 model_mappings 配置的）

        Args:
            db: 数据库 Session
            provider_id: Provider ID

        Returns:
            Dict: 包含 success 和 errors 列表
        """
        from src.core.model_permissions import match_model_with_pattern
        from src.models.database import Provider, ProviderAPIKey

        results: dict[str, list[dict]] = {
            "success": [],
            "errors": [],
        }

        # 获取 Provider
        provider = db.query(Provider).filter(Provider.id == provider_id).first()
        if not provider:
            logger.warning(f"Provider {provider_id} not found for auto-disassociation")
            return results

        # 1. 先快速检查是否存在"允许所有模型"的活跃 Key。
        # 这种情况下无需解除任何关联，避免继续扫描整张 key 表。
        # 注意：跳过 OAuth Key，OAuth Key 的 allowed_models 由上游动态获取，数量庞大，
        # 不应参与 disassociate 判定。
        from src.services.provider_keys.auth_type import OAUTH_AUTH_TYPES

        non_oauth_filter = ProviderAPIKey.auth_type.notin_(OAUTH_AUTH_TYPES)
        has_unlimited_key = (
            db.query(ProviderAPIKey.id)
            .filter(
                ProviderAPIKey.provider_id == provider_id,
                ProviderAPIKey.is_active == True,
                ProviderAPIKey.allowed_models.is_(None),
                non_oauth_filter,
            )
            .limit(1)
            .first()
            is not None
        )
        if has_unlimited_key:
            return results

        # 2. 仅查询活跃 Key 的 allowed_models 列，避免把 api_key/auth_config 等大字段整行拉出。
        allowed_model_rows = (
            db.query(ProviderAPIKey.allowed_models)
            .filter(
                ProviderAPIKey.provider_id == provider_id,
                ProviderAPIKey.is_active == True,
                non_oauth_filter,
            )
            .all()
        )

        # 如果 Provider 无活跃 Key，不做任何解除（保留现有关联）
        if not allowed_model_rows:
            return results

        # 收集所有 Key 的 allowed_models 并集
        all_allowed_models: set[str] = set()
        for (allowed_models,) in allowed_model_rows:
            if isinstance(allowed_models, list) and allowed_models:
                all_allowed_models.update(m for m in allowed_models if isinstance(m, str))

        # 3. 获取 Provider 当前关联的所有 Model（仅加载判定所需字段）
        models = (
            db.query(Model)
            .options(
                load_only(Model.id, Model.provider_id, Model.global_model_id),
                joinedload(Model.global_model).load_only(
                    GlobalModel.id,
                    GlobalModel.name,
                    GlobalModel.config,
                ),
            )
            .filter(Model.provider_id == provider_id)
            .all()
        )

        # 4. 检查每个 Model 是否还能匹配，收集需要删除的 Model
        models_to_delete: list[Model] = []

        for model in models:
            # 跳过 global_model 关系未加载的
            if not model.global_model:
                continue

            global_model = cast(GlobalModel, model.global_model)

            # 提取映射规则
            model_mappings: list[str] = []
            config = global_model.config
            if config and isinstance(config, dict):
                mappings = config.get("model_mappings")
                if isinstance(mappings, list):
                    model_mappings = [m for m in mappings if isinstance(m, str)]

            # 如果 GlobalModel 没有 model_mappings，跳过（说明不是通过映射自动关联的）
            if not model_mappings:
                continue

            # 检查是否有映射规则匹配到任一 allowed_models
            matched = False
            for mapping_pattern in model_mappings:
                for allowed_model in all_allowed_models:
                    if match_model_with_pattern(mapping_pattern, allowed_model):
                        matched = True
                        break
                if matched:
                    break

            # 如果不再匹配，标记为待删除
            if not matched:
                models_to_delete.append(model)

        # 5. 批量删除不再匹配的 Model（全部成功或全部失败）
        if models_to_delete:
            try:
                for model in models_to_delete:
                    global_model = cast(GlobalModel, model.global_model)
                    db.delete(model)
                    results["success"].append(
                        {
                            "model_id": model.id,
                            "global_model_id": global_model.id,
                            "global_model_name": global_model.name,
                        }
                    )
                    logger.info(
                        f"Auto-disassociated Provider {provider.name} from GlobalModel {global_model.name} "
                        f"(no matching allowed_models)"
                    )
                db.commit()
            except Exception as e:
                db.rollback()
                logger.error(f"Failed to auto-disassociate Provider {provider.name}: {e}")
                # 清空 success，记录整体错误
                results["success"] = []
                results["errors"].append(
                    {
                        "provider_id": provider_id,
                        "error": str(e),
                    }
                )

        return results