"""
LLM API客户端工具类
支持异步并发调用多个LLM API请求
"""

from __future__ import annotations
import asyncio
import json
from typing import Any, Dict, List, Optional
import re

try:
    import aiohttp
    HAS_AIOHTTP = True
except ImportError:
    HAS_AIOHTTP = False

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

from ..config.provider import default_config_provider
from foundation.infrastructure.config.config import config_handler
from foundation.observability.logger.loggering import review_logger as logger


class LLMClient:
    """LLM API客户端，支持异步并发调用"""

    def __init__(self, config_provider=None):
        """
        初始化LLM客户端
        
        参数:
            config_provider: 配置提供者，如果为None则使用默认配置
        """
        self._cfg = config_provider or default_config_provider
        self._load_config()

    def _load_config(self):
        """加载LLM API配置（优先从 model_setting.yaml，回退到 config.ini）"""
        # 获取模型类型（优先从 model_setting.yaml 读取默认配置）
        try:
            from foundation.ai.models.model_config_loader import get_model_for_function
            model_type = get_model_for_function("default")
            if model_type:
                self.model_type = model_type.lower()
                logger.debug(f"LLMClient 从 model_setting.yaml 读取默认模型: {self.model_type}")
            else:
                self.model_type = config_handler.get("model", "MODEL_TYPE", "qwen3_5_35b_a3b").lower()
        except Exception as e:
            logger.debug(f"LLMClient 从 model_setting.yaml 读取失败: {e}，回退到 config.ini")
            self.model_type = config_handler.get("model", "MODEL_TYPE", "qwen3_5_35b_a3b").lower()

        # 获取模型配置（根据模型类型动态读取对应节）
        server_url = ""
        model_id = ""
        api_key = ""

        if self.model_type.startswith("shutian"):
            # 蜀天模型系列 - 从 shutian 节读取配置
            # 注意：判断顺序很重要，122b 包含 35b 字符串，要先判断 122b
            if "122b" in self.model_type:
                server_url = config_handler.get("shutian", "SHUTIAN_122B_SERVER_URL", "")
                model_id = config_handler.get("shutian", "SHUTIAN_122B_MODEL_ID", "")
                api_key = config_handler.get("shutian", "SHUTIAN_122B_API_KEY", "")
            elif "8b" in self.model_type and "35b" not in self.model_type and "122b" not in self.model_type:
                server_url = config_handler.get("shutian", "SHUTIAN_8B_SERVER_URL", "")
                model_id = config_handler.get("shutian", "SHUTIAN_8B_MODEL_ID", "")
                api_key = config_handler.get("shutian", "SHUTIAN_8B_API_KEY", "")
            elif "35b" in self.model_type:
                # 35B 模型
                server_url = config_handler.get("shutian", "SHUTIAN_35B_SERVER_URL", "")
                model_id = config_handler.get("shutian", "SHUTIAN_35B_MODEL_ID", "")
                api_key = config_handler.get("shutian", "SHUTIAN_35B_API_KEY", "")
            else:
                # 默认 35B 模型
                server_url = config_handler.get("shutian", "SHUTIAN_35B_SERVER_URL", "")
                model_id = config_handler.get("shutian", "SHUTIAN_35B_MODEL_ID", "")
                api_key = config_handler.get("shutian", "SHUTIAN_35B_API_KEY", "")
        else:
            # 其他模型 - 从对应节读取 DashScope 配置
            server_url = config_handler.get(self.model_type, "DASHSCOPE_SERVER_URL", "")
            model_id = config_handler.get(self.model_type, "DASHSCOPE_MODEL_ID", "")
            api_key = config_handler.get(self.model_type, "DASHSCOPE_API_KEY", "")

            # 如果 DashScope 配置不存在，尝试读取其他模型配置（兼容旧配置）
            if not server_url:
                # 尝试读取 QWEN_SERVER_URL 等旧格式配置
                server_url = config_handler.get(self.model_type, f"{self.model_type.upper()}_SERVER_URL", "")
                model_id = config_handler.get(self.model_type, f"{self.model_type.upper()}_MODEL_ID", "")
                api_key = config_handler.get(self.model_type, f"{self.model_type.upper()}_API_KEY", "")

        self.api_url = server_url.rstrip("/") if server_url else ""
        self.model_id = model_id
        self.api_key = api_key
        self.base_url = f"{self.api_url}/chat/completions" if self.api_url else ""

        # 通用配置
        self.timeout = int(config_handler.get("llm_keywords", "TIMEOUT", "60"))
        self.max_retries = int(config_handler.get("llm_keywords", "MAX_RETRIES", "2"))
        self.concurrent_workers = int(config_handler.get("llm_keywords", "CONCURRENT_WORKERS", "20"))
        self.stream = config_handler.get("llm_keywords", "STREAM", "false").lower() == "true"
        self.temperature = float(config_handler.get("llm_keywords", "TEMPERATURE", "0.3"))
        self.max_tokens = int(config_handler.get("llm_keywords", "MAX_TOKENS", "1024"))

    def _extract_json_from_string(self, text: str) -> Optional[Dict[str, Any]]:
        """
        从字符串中提取第一个有效的JSON对象。
        尝试处理JSON被markdown代码块包裹的情况。
        """
        # 1. 尝试从 ```json ... ``` 代码块中提取
        match = re.search(r"```json\s*(\{.*?})\s*```", text, re.DOTALL)
        if match:
            json_str = match.group(1)
            try:
                return json.loads(json_str)
            except json.JSONDecodeError:
                pass # 继续尝试其他方式

        # 2. 尝试从 ``` ... ``` 代码块中提取
        match = re.search(r"```\s*(\{.*?})\s*```", text, re.DOTALL)
        if match:
            json_str = match.group(1)
            try:
                return json.loads(json_str)
            except json.JSONDecodeError:
                pass # 继续尝试其他方式
        
        # 3. 尝试直接从字符串中查找第一个JSON对象
        # 寻找第一个 { 和最后一个 }
        try:
            # 查找所有可能的JSON对象
            json_objects = re.findall(r"(\{.*?\})", text, re.DOTALL)
            for json_str in json_objects:
                try:
                    return json.loads(json_str)
                except json.JSONDecodeError:
                    pass
        except Exception:
            pass

        return None

    async def _call_api_async(self, session: aiohttp.ClientSession, messages: List[Dict[str, str]]) -> Dict[str, Any]:
        """
        异步调用LLM API
        
        参数:
            session: aiohttp会话
            messages: 消息列表
            
        返回:
            API响应结果
        """
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        payload = {
            "model": self.model_id,
            "messages": messages,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "stream": self.stream
        }
        
        for attempt in range(self.max_retries):
            try:
                async with session.post(
                    self.base_url,
                    json=payload,
                    headers=headers,
                    timeout=aiohttp.ClientTimeout(total=self.timeout)
                ) as response:
                    if response.status == 200:
                        result = await response.json()
                        return result
                    else:
                        error_text = await response.text()
                        if attempt < self.max_retries - 1:
                            await asyncio.sleep(1 * (attempt + 1))  # 指数退避
                            continue
                        raise Exception(f"API调用失败，状态码: {response.status}, 错误: {error_text}")
            except asyncio.TimeoutError:
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(1 * (attempt + 1))
                    continue
                raise Exception(f"API调用超时（超过{self.timeout}秒）")
            except Exception as e:
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(1 * (attempt + 1))
                    continue
                raise
        
        raise Exception("API调用失败，已达到最大重试次数")

    def _call_api_sync(self, messages: List[Dict[str, str]]) -> Dict[str, Any]:
        """
        同步调用LLM API（回退方案，当没有aiohttp时使用）
        
        参数:
            messages: 消息列表
            
        返回:
            API响应结果
        """
        if not HAS_REQUESTS:
            raise ImportError("需要安装 aiohttp 或 requests 库才能使用LLM API客户端")
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        payload = {
            "model": self.model_id,
            "messages": messages,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "stream": self.stream
        }
        
        for attempt in range(self.max_retries):
            try:
                response = requests.post(
                    self.base_url,
                    json=payload,
                    headers=headers,
                    timeout=self.timeout
                )
                if response.status_code == 200:
                    return response.json()
                else:
                    if attempt < self.max_retries - 1:
                        import time
                        time.sleep(1 * (attempt + 1))
                        continue
                    raise Exception(f"API调用失败，状态码: {response.status_code}, 错误: {response.text}")
            except requests.Timeout:
                if attempt < self.max_retries - 1:
                    import time
                    time.sleep(1 * (attempt + 1))
                    continue
                raise Exception(f"API调用超时（超过{self.timeout}秒）")
            except Exception as e:
                if attempt < self.max_retries - 1:
                    import time
                    time.sleep(1 * (attempt + 1))
                    continue
                raise
        
        raise Exception("API调用失败，已达到最大重试次数")

    async def _process_single_request(self, session: aiohttp.ClientSession, messages: List[Dict[str, str]]) -> Optional[Dict[str, Any]]:
        """
        处理单个请求（包装异常处理）
        
        参数:
            session: aiohttp会话
            messages: 消息列表
            
        返回:
            解析后的JSON结果，如果失败则返回None
        """
        try:
            response = await self._call_api_async(session, messages)
            
            # 提取响应内容
            if "choices" in response and len(response["choices"]) > 0:
                content = response["choices"][0].get("message", {}).get("content", "")
                
                # 尝试解析JSON
                extracted_json = self._extract_json_from_string(content)
                if extracted_json:
                    return extracted_json
                else:
                    # 如果不是JSON，返回原始内容
                    return {"raw_content": content}
            else:
                return None
        except Exception as e:
            logger.error(f"[LLMClient] LLM API调用错误: {e}")
            return None

    async def batch_call_async(self, requests: List[List[Dict[str, str]]]) -> List[Optional[Dict[str, Any]]]:
        """
        异步批量调用LLM API
        
        参数:
            requests: 请求列表，每个请求是一个消息列表
            
        返回:
            结果列表，与输入请求一一对应
        """
        if not HAS_AIOHTTP:
            # 回退到同步调用（在异步环境中）
            if HAS_REQUESTS:
                logger.warning("[LLMClient] 未安装aiohttp，在异步环境中使用同步调用（性能较差）")
                results = []
                for req in requests:
                    try:
                        response = self._call_api_sync(req)
                        if "choices" in response and len(response["choices"]) > 0:
                            content = response["choices"][0].get("message", {}).get("content", "")
                            try:
                                if "```json" in content:
                                    start = content.find("```json") + 7
                                    end = content.find("```", start)
                                    content = content[start:end].strip()
                                elif "```" in content:
                                    start = content.find("```") + 3
                                    end = content.find("```", start)
                                    content = content[start:end].strip()
                                results.append(json.loads(content))
                            except json.JSONDecodeError:
                                results.append({"raw_content": content})
                        else:
                            results.append(None)
                    except Exception as e:
                        logger.error(f"[LLMClient] LLM API调用错误: {e}")
                        results.append(None)
                return results
            else:
                raise ImportError("需要安装 aiohttp 或 requests 库才能使用LLM API客户端")
        
        # 使用信号量限制并发数
        semaphore = asyncio.Semaphore(self.concurrent_workers)
        
        async def bounded_request(session, messages):
            async with semaphore:
                return await self._process_single_request(session, messages)
        
        async with aiohttp.ClientSession() as session:
            tasks = [bounded_request(session, req) for req in requests]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 处理异常结果
            processed_results = []
            for result in results:
                if isinstance(result, Exception):
                    logger.error(f"[LLMClient] LLM API调用异常: {result}")
                    processed_results.append(None)
                else:
                    processed_results.append(result)
            
            return processed_results

    def batch_call(self, requests: List[List[Dict[str, str]]]) -> List[Optional[Dict[str, Any]]]:
        """
        同步批量调用LLM API（兼容接口）

        注意: 此方法使用 workflow_manager.py 的全局事件循环，不再自行初始化事件循环
        """
        if HAS_AIOHTTP:
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    logger.warning("[LLMClient] 检测到运行中的事件循环，batch_call 请改用 await batch_call_async；本次回退同步调用")
                    return self._batch_call_sync_fallback(requests)
                logger.debug("[LLMClient] 异步调用LLM API进行目录分类处理")
                return loop.run_until_complete(self.batch_call_async(requests))
            except RuntimeError:
                logger.debug("[LLMClient] 同步调用LLM API进行目录分类处理（无事件循环）")
                return self._batch_call_sync_fallback(requests)
        else:
            return self._batch_call_sync_fallback(requests)

    def _batch_call_sync_fallback(self, requests: List[List[Dict[str, str]]]) -> List[Optional[Dict[str, Any]]]:
        """
        同步批量调用回退方案
        """
        if not HAS_REQUESTS:
            raise ImportError("需要安装 requests 库才能使用同步调用模式")
        
        results = []
        for req in requests:
            try:
                response = self._call_api_sync(req)
                if "choices" in response and len(response["choices"]) > 0:
                    content = response["choices"][0].get("message", {}).get("content", "")
                    try:
                        extracted_json = self._extract_json_from_string(content)
                        if extracted_json:
                            results.append(extracted_json)
                        else:
                            results.append({"raw_content": content})
                    except Exception:
                        results.append({"raw_content": content})
                else:
                    results.append(None)
            except Exception as e:
                logger.error(f"[LLMClient] LLM API调用错误: {e}")
                results.append(None)
        return results

