src/model_tester.py

import asyncio
import aiohttp
from typing import List, Dict, Any, Optional, Set
from datetime import datetime
import random

from .openrouter_client import OpenRouterClient
from . import config
from .utils import clean_model_name


def _is_free_model(model: dict) -> bool:
    """判断模型是否免费（:free 后缀 or pricing 全为 0）"""
    model_id = model.get("id", "")
    if ":free" in model_id:
        return True
    pricing = model.get("pricing", {}) or {}
    prompt_price = pricing.get("prompt", "0")
    completion_price = pricing.get("completion", "0")
    request_price = pricing.get("request", "0")
    return prompt_price == "0" and completion_price == "0" and request_price == "0"


class ModelTester:
    def __init__(self):
        self.client = OpenRouterClient()
        self.max_concurrency = config.get_max_concurrency()
        self.test_prompt = config.get_test_prompt()

        self._all_models: List[str] = []
        self._free_models: List[str] = []
        self._available_models: List[str] = []
        self._available_free_models: List[str] = []
        self._scan_in_progress = False
        self._last_scan_time: Optional[datetime] = None
        self.scan_result: Dict[str, Any] = {
            "available_models": [],
            "available_free_models": [],
            "total_available": 0,
            "free_available": 0,
            "timestamp": None
        }

    # ------------------------------------------------------------------ #
    #  模型列表刷新
    # ------------------------------------------------------------------ #

    def refresh_model_list(self):
        """Get latest model list from API (sync, legacy)"""
        models = self.client.get_models()
        all_ids = []
        free_ids = []

        for model in models:
            model_id = model.get("id", "")
            if model_id:
                all_ids.append(model_id)
                if ":free" in model_id:
                    free_ids.append(model_id)

        self._all_models = all_ids
        self._free_models = free_ids
        return len(self._all_models), len(self._free_models)

    async def refresh_model_list_async(self):
        """Get latest model list from API (async, non-blocking)"""
        models = await self.client.async_get_models()
        all_ids = []
        free_ids = []

        for model in models:
            model_id = model.get("id", "")
            if model_id:
                all_ids.append(model_id)
                if _is_free_model(model):
                    free_ids.append(model_id)

        self._all_models = all_ids
        self._free_models = free_ids
        return len(self._all_models), len(self._free_models)

    # ------------------------------------------------------------------ #
    #  模型全量扫描（定时任务用）
    # ------------------------------------------------------------------ #

    async def test_single_model_async(
        self,
        session: aiohttp.ClientSession,
        model_id: str,
        api_key: str
    ) -> tuple[str, bool, str]:
        """测试单个模型，返回 (model_id, success, api_key)"""
        url = "https://openrouter.ai/api/v1/chat/completions"
        payload = {
            "model": model_id,
            "messages": [{"role": "user", "content": self.test_prompt}],
            "max_tokens": 10
        }
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

        try:
            timeout = aiohttp.ClientTimeout(total=config.get_request_timeout())
            async with session.post(url, json=payload, headers=headers, timeout=timeout) as response:
                is_success = response.status == 200
                return model_id, is_success, api_key
        except Exception:
            return model_id, False, api_key

    async def scan_all_models_async(self):
        """Async scan all models concurrently with multi-key distribution"""
        if self._scan_in_progress:
            return {"error": "Scan already in progress"}

        self._scan_in_progress = True
        print(f"[{datetime.now()}] Starting model scan...")

        all_count, free_count = await self.refresh_model_list_async()
        print(f"Total models: {all_count}, Free models: {free_count}")

        api_keys = config.get_api_keys()

        available: Set[str] = set()
        available_free: Set[str] = set()

        async with aiohttp.ClientSession() as session:
            semaphore = asyncio.Semaphore(self.max_concurrency)

            async def test_with_semaphore(model_id: str, api_key: str):
                async with semaphore:
                    return await self.test_single_model_async(session, model_id, api_key)

            tasks = [
                test_with_semaphore(m, api_keys[i % len(api_keys)])
                for i, m in enumerate(self._all_models)
            ]
            results = await asyncio.gather(*tasks, return_exceptions=True)

            for result in results:
                if isinstance(result, tuple):
                    model_id, success, _ = result
                    cleaned = clean_model_name(model_id)
                    if success:
                        available.add(cleaned)
                        if ":free" in model_id:
                            available_free.add(cleaned)

        self._available_models = sorted(list(available))
        self._available_free_models = sorted(list(available_free))
        self._last_scan_time = datetime.now()
        self._scan_in_progress = False

        self.scan_result = {
            "available_models": self._available_models,
            "available_free_models": self._available_free_models,
            "total_available": len(self._available_models),
            "free_available": len(self._available_free_models),
            "timestamp": self._last_scan_time.isoformat() if self._last_scan_time else None
        }

        print(f"Scan complete: {len(self._available_free_models)} free, {len(self._available_models)} total available")

        return self.scan_result

    def scan_all_models(self):
        """Sync wrapper for scan"""
        return asyncio.run(self.scan_all_models_async())

    def get_available_models(self, free_only: bool = False) -> List[str]:
        if free_only:
            return self._available_free_models
        return self._available_models

    def get_all_free_models(self) -> List[str]:
        return self._free_models

    # ------------------------------------------------------------------ #
    #  单模型直接请求
    # ------------------------------------------------------------------ #

    async def try_model_direct(
        self,
        session: aiohttp.ClientSession,
        model_id: str,
        api_key: str,
        prompt: str = None
    ) -> Optional[Dict[str, Any]]:
        url = "https://openrouter.ai/api/v1/chat/completions"
        payload = {
            "model": model_id,
            "messages": [{"role": "user", "content": prompt or self.test_prompt}]
        }
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

        try:
            timeout = aiohttp.ClientTimeout(total=config.get_request_timeout())
            async with session.post(url, json=payload, headers=headers, timeout=timeout) as response:
                if response.status == 200:
                    data = await response.json()
                    return {
                        "success": True,
                        "model": model_id,
                        "response": data,
                        "method": "direct"
                    }
                else:
                    body = await response.text()
                    print(f"[try_model_direct] ERROR {model_id}: HTTP {response.status}, body: {body[:200]}")
                    return {
                        "success": False,
                        "model": model_id,
                        "error": f"HTTP {response.status}: {body[:100]}",
                        "method": "direct"
                    }
        except asyncio.TimeoutError:
            return {"success": False, "model": model_id, "error": "timeout", "method": "direct"}
        except Exception as e:
            return {"success": False, "model": model_id, "error": str(e), "method": "direct"}

    async def try_model_direct_stream(
        self,
        session: aiohttp.ClientSession,
        model_id: str,
        api_key: str,
        messages: List[Dict[str, str]]
    ):
        """发送流式请求到OpenRouter，返回流式迭代器"""
        url = "https://openrouter.ai/api/v1/chat/completions"
        payload = {
            "model": model_id,
            "messages": messages,
            "stream": True
        }
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

        async with session.post(url, json=payload, headers=headers, timeout=aiohttp.ClientTimeout(total=120)) as response:
            async for line in response.content:
                yield line

    # ------------------------------------------------------------------ #
    #  工具方法：模型匹配
    # ------------------------------------------------------------------ #

    @staticmethod
    def _model_name(model_id: str) -> str:
        """从完整 model ID 中提取模型名，如 'openai/gpt-4o:free' → 'gpt-4o'"""
        return model_id.replace(":free", "").split("/")[-1]

    def _match_free_models(self, keyword: str) -> List[str]:
        """Tier 1: 从带 :free 标签的模型列表中按关键词匹配"""
        matched = []
        for m in self._free_models:
            if keyword.lower() in self._model_name(m).lower():
                matched.append(m)
        return matched[:10]

    def _match_paid_models(self, keyword: str) -> List[str]:
        """Tier 2: 从全部模型列表中匹配非 free 的模型"""
        free_set = set(self._free_models)
        matched = []
        for m in self._all_models:
            if m in free_set:  # 跳过 free 模型
                continue
            if keyword.lower() in self._model_name(m).lower():
                matched.append(m)
        return matched[:10]

    def _match_available_models(self, keyword: str) -> List[str]:
        """Tier 3: 从已验证可用的模型列表中模糊匹配（已清洗，无前缀无:free）"""
        matched = []
        for m in self._available_models:
            if keyword.lower() in m.lower():
                matched.append(m)
        return matched[:10]

    # ------------------------------------------------------------------ #
    #  三级并发尝试
    # ------------------------------------------------------------------ #

    async def _call_with_retry(
        self,
        session: aiohttp.ClientSession,
        model_id: str,
        api_keys: List[str],
        prompt: str = None
    ) -> Optional[Dict[str, Any]]:
        """
        调用 try_model_direct 并附带重试逻辑。

        报错（timeout / 连接错误 / 429 / 5xx）→ 换 key + 退避重试
        未找到（400 / 404 / 403 / 401）      → 不重试，直接返回失败（触发降级）
        """
        failed_keys: set = set()
        result = None
        max_retries = config.get_max_retries()
        retry_delay = config.get_retry_delay()

        for attempt in range(max_retries):
            available = [k for k in api_keys if k not in failed_keys]
            if not available:
                print(f"[_call_with_retry] All keys failed 401, giving up on {model_id}")
                break
            key = random.choice(available)

            result = await self.try_model_direct(session, model_id, key, prompt)

            if result is None:
                continue
            if result.get("success"):
                return result

            error = str(result.get("error", ""))

            # ── 未找到 → 不重试，立即降级 ──
            if any(s in error for s in ("HTTP 400", "HTTP 403", "HTTP 404", "model_not_found", "not found")):
                print(f"[_call_with_retry] {model_id} not found (definitive), skip retry")
                return result

            # ── 401 → 换 key，不等待 ──
            if "401" in error or "unauthorized" in error.lower():
                print(f"[_call_with_retry] {model_id} 401, rotating key")
                failed_keys.add(key)
                continue

            # ── 报错 → 退避重试 ──
            sleep_s = retry_delay * (attempt + 1)
            print(f"[_call_with_retry] {model_id} transient error ({error}), retry {attempt + 1}/{max_retries} in {sleep_s}s")
            await asyncio.sleep(sleep_s)

        return result

    async def _try_candidates_concurrently(
        self,
        session: aiohttp.ClientSession,
        candidates: List[str],
        api_keys: List[str],
        prompt: str,
        label: str
    ) -> Optional[Dict[str, Any]]:
        """并发测试一组候选模型（含重试），返回第一个成功的"""
        if not candidates:
            return None

        semaphore = asyncio.Semaphore(min(5, len(candidates)))

        async def try_one(model_id: str):
            async with semaphore:
                print(f"[{label}] Testing: {model_id}")
                result = await self._call_with_retry(session, model_id, api_keys, prompt)
                if result and result.get("success"):
                    print(f"[{label}] SUCCESS: {model_id}")
                    return result
                return None

        tasks = [try_one(m) for m in candidates]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        for r in results:
            if isinstance(r, dict) and r.get("success"):
                r["method"] = label
                return r
        return None

    # ------------------------------------------------------------------ #
    #  三级模型选择：chat_completion（非流式）
    # ------------------------------------------------------------------ #

    async def chat_completion(self, prompt: str, model_hint: Optional[str] = None) -> Dict[str, Any]:
        """
        三级模型选择策略：
          1. 从 free 模型列表中匹配 hint → 并发尝试
          2. 从非 free（付费）模型列表中匹配 hint → 并发尝试
          3. 从已验证可用的列表中模糊匹配 hint → 并发尝试
        """
        api_keys = config.get_api_keys()

        await self.refresh_model_list_async()

        async with aiohttp.ClientSession() as session:
            if not model_hint:
                # 没有提示词，直接用已知可用的 free 模型
                result = await self._try_candidates_concurrently(
                    session, self._free_models[:15], api_keys, prompt, "random"
                )
                if result and result.get("success"):
                    return {
                        "success": True,
                        "response": result.get("response"),
                        "method": result.get("method"),
                        "model": result.get("model")
                    }
                return {
                    "success": False,
                    "error": "No available model",
                    "method": "no_candidates"
                }

            # T1: free 模型
            candidates = self._match_free_models(model_hint)
            if candidates:
                result = await self._try_candidates_concurrently(session, candidates, api_keys, prompt, "t1_free")
                if result:
                    return {
                        "success": True,
                        "response": result.get("response"),
                        "method": result.get("method"),
                        "model": result.get("model")
                    }

            # T2: 付费模型
            candidates = self._match_paid_models(model_hint)
            if candidates:
                result = await self._try_candidates_concurrently(session, candidates, api_keys, prompt, "t2_paid")
                if result:
                    return {
                        "success": True,
                        "response": result.get("response"),
                        "method": result.get("method"),
                        "model": result.get("model")
                    }

            # T3: 已测试列表模糊匹配
            candidates = self._match_available_models(model_hint)
            if candidates:
                result = await self._try_candidates_concurrently(session, candidates, api_keys, prompt, "t3_tested")
                if result:
                    return {
                        "success": True,
                        "response": result.get("response"),
                        "method": result.get("method"),
                        "model": result.get("model")
                    }

            return {
                "success": False,
                "error": f"No available model for '{model_hint}'",
                "method": "all_tiers_failed"
            }

    def chat_completion_sync(self, prompt: str, model_hint: Optional[str] = None) -> Dict[str, Any]:
        return asyncio.run(self.chat_completion(prompt, model_hint))

    # ------------------------------------------------------------------ #
    #  三级模型选择：chat_completion_stream（流式）
    # ------------------------------------------------------------------ #

    async def chat_completion_stream(self, model_hint: Optional[str], messages: List[Dict[str, str]]):
        """
        流式聊天，同样三级策略（流式必须逐个尝试，没法并发）。
        """
        api_keys = config.get_api_keys()
        api_key = random.choice(api_keys)

        await self.refresh_model_list_async()

        if not model_hint:
            # 无 hint，直接尝试已知 free 模型
            async with aiohttp.ClientSession() as session:
                for model_id in self._free_models[:10]:
                    try:
                        async for chunk in self.try_model_direct_stream(session, model_id, api_key, messages):
                            yield chunk
                        return
                    except Exception as e:
                        print(f"[stream] Free model {model_id} failed: {e}")
                        continue
            return

        # T1: free 模型
        candidates = self._match_free_models(model_hint)
        if candidates:
            async with aiohttp.ClientSession() as session:
                for model_id in candidates:
                    try:
                        async for chunk in self.try_model_direct_stream(session, model_id, api_key, messages):
                            yield chunk
                        return
                    except Exception:
                        continue

        # T2: 付费模型
        candidates = self._match_paid_models(model_hint)
        if candidates:
            async with aiohttp.ClientSession() as session:
                for model_id in candidates:
                    try:
                        async for chunk in self.try_model_direct_stream(session, model_id, api_key, messages):
                            yield chunk
                        return
                    except Exception:
                        continue

        # T3: 已测试列表模糊匹配
        candidates = self._match_available_models(model_hint)
        if candidates:
            async with aiohttp.ClientSession() as session:
                for model_id in candidates:
                    try:
                        async for chunk in self.try_model_direct_stream(session, model_id, api_key, messages):
                            yield chunk
                        return
                    except Exception:
                        continue

    # ------------------------------------------------------------------ #
    #  兼容旧接口
    # ------------------------------------------------------------------ #

    def test_single_model(self, model_id: str) -> tuple[str, bool]:
        is_available = self.client.test_model(model_id, self.test_prompt)
        cleaned_name = clean_model_name(model_id)
        return cleaned_name, is_available

    def test_all_models(self) -> Dict[str, Any]:
        return self.scan_all_models()