# SPDX-FileCopyrightText: Copyright contributors to the RouterArena project
# SPDX-License-Identifier: Apache-2.0
#
# 本文件 vendored 自 RouteWorks/RouterArena（https://github.com/RouteWorks/RouterArena，
# 2026-08-19 拉取，commit on main），仅保留本项目适配所需最小接口定义。
# 原始完整代码与本项目无关；如 RouterArena 接口变化，需同步更新本文件。

"""RouterArena BaseRouter 最小 vendored 实现（适配本系统需要）。

仅保留以下能力：
- 配置加载与模型列表提取
- 模型名校验
- 抽象方法 _get_prediction

去掉了原项目对 generate_prediction_file 路径的硬编码（用 ConfigResolver 解耦）。
"""
from __future__ import annotations

import json
import os
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional


class BaseRouter(ABC):
    """Abstract base class for router implementations.

    子类必须实现 _get_prediction(query) -> str，返回 config.models 中存在的模型名。
    """

    def __init__(
        self,
        router_name: str,
        config_path: Optional[str] = None,
    ):
        self.router_name = router_name
        # 允许外部注入 config_path，避免 RouterArena 仓库的硬编码路径依赖
        if config_path is None:
            config_path = self._default_config_path()
        self.config_path = config_path
        self.config = self._load_config()
        self.models = self._extract_models()

    def _default_config_path(self) -> str:
        # 默认查找约定：<project>/router_inference/config/<router_name>.json
        # 优先尝试项目内 research/routerarena/config/，再退到 RouterArena 约定路径
        here = os.path.dirname(os.path.abspath(__file__))
        candidate = os.path.join(here, "config", f"{self.router_name}.json")
        if os.path.exists(candidate):
            return candidate
        return candidate  # 不存在时让 _load_config 抛 FileNotFoundError，给出明确路径

    def _load_config(self) -> Dict[str, Any]:
        if not os.path.exists(self.config_path):
            raise FileNotFoundError(f"Config file not found: {self.config_path}")
        with open(self.config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
        if "pipeline_params" not in config:
            raise ValueError(f"Invalid config: missing 'pipeline_params' in {self.config_path}")
        if "models" not in config["pipeline_params"]:
            raise ValueError(f"Invalid config: missing 'models' in pipeline_params")
        return config

    def _extract_models(self) -> List[str]:
        return list(self.config["pipeline_params"]["models"])

    def _validate_model(self, model_name: str) -> None:
        if model_name not in self.models:
            raise ValueError(
                f"Model '{model_name}' not in router config. "
                f"Available: {self.models}"
            )

    @abstractmethod
    def _get_prediction(self, query: str) -> str:
        """根据 query 返回 config.models 中存在的目标模型名。"""
        raise NotImplementedError

    def get_prediction(self, query: str) -> str:
        model_name = self._get_prediction(query)
        self._validate_model(model_name)
        return model_name