# Copyright 2023 The Qwen team, Alibaba Group. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import json
import logging
import os
import time
from typing import Dict, Iterator, List, Literal, Optional, Union

from qwen_agent.agents import Assistant
from qwen_agent.llm.schema import ASSISTANT, FUNCTION, Message
from qwen_agent.llm.oai import TextChatAtOAI 
from qwen_agent.tools import BaseTool
from agent.custom_mcp_manager import CustomMCPManager
import logging

logger = logging.getLogger('app')
# 设置工具日志记录器
tool_logger = logging.getLogger('app')

class ModifiedAssistant(Assistant):
    """
    修改后的 Assistant 子类，改变循环判断逻辑：
    - 原始逻辑：如果没有使用工具，立即退出循环
    - 修改后逻辑：如果没有使用工具，调用模型判断回答是否完整，如果不完整则继续循环
    """
    
    def _is_retryable_error(self, error: Exception) -> bool:
        """判断错误是否可重试
        
        Args:
            error: 异常对象
            
        Returns:
            bool: 是否可重试
        """
        error_str = str(error).lower()
        retryable_indicators = [
            '502', '500', '503', '504',  # HTTP错误代码
            'internal server error',     # 内部服务器错误
            'timeout',                   # 超时
            'connection',                # 连接错误
            'network',                   # 网络错误
            'rate',                      # 速率限制和相关错误
            'quota',                     # 配额限制
            'service unavailable',      # 服务不可用
            'provider returned error',   # Provider错误
            'model service error',       # 模型服务错误
            'temporary',                # 临时错误
            'retry'                     # 明确提示重试
        ]
        return any(indicator in error_str for indicator in retryable_indicators)
    
    def _call_tool(self, tool_name: str, tool_args: Union[str, dict] = '{}', **kwargs) -> str:
        """重写工具调用方法，添加调试信息"""
        if tool_name not in self.function_map:
            error_msg = f'Tool {tool_name} does not exist. Available tools: {list(self.function_map.keys())}'
            tool_logger.error(error_msg)
            return error_msg

        tool = self.function_map[tool_name]

        try:
            tool_logger.info(f"开始调用工具: {tool_name} {tool_args}")
            start_time = time.time()

            # 调用父类的_call_tool方法
            tool_result = super()._call_tool(tool_name, tool_args, **kwargs)

            end_time = time.time()
            tool_logger.info(f"工具 {tool_name} 执行完成，耗时: {end_time - start_time:.2f}秒 结果长度: {len(tool_result) if tool_result else 0}")

            # 打印部分结果内容（避免过长）
            if tool_result and len(tool_result) > 0:
                preview = tool_result[:200] if len(tool_result) > 200 else tool_result
                tool_logger.debug(f"工具 {tool_name} 结果预览: {preview}...")

            return tool_result

        except Exception as ex:
            end_time = time.time()
            tool_logger.error(f"工具调用异常，耗时: {end_time - start_time:.2f}秒 异常类型: {type(ex).__name__} {str(ex)}")

            # 打印完整的堆栈跟踪
            import traceback
            tool_logger.error(f"堆栈跟踪:\n{traceback.format_exc()}")

            # 返回详细的错误信息
            error_message = f'An error occurred when calling tool {tool_name}: {type(ex).__name__}: {str(ex)}'
            return error_message

    def _init_tool(self, tool: Union[str, Dict, BaseTool]):
        """重写工具初始化方法，使用CustomMCPManager处理MCP服务器配置"""
        if isinstance(tool, BaseTool):
            # 处理BaseTool实例
            tool_name = tool.name
            if tool_name in self.function_map:
                tool_logger.warning(f'Repeatedly adding tool {tool_name}, will use the newest tool in function list')
            self.function_map[tool_name] = tool
        elif isinstance(tool, dict) and 'mcpServers' in tool:
            # 使用CustomMCPManager处理MCP服务器配置，支持headers
            tools = CustomMCPManager().initConfig(tool)
            for tool in tools:
                tool_name = tool.name
                if tool_name in self.function_map:
                    tool_logger.warning(f'Repeatedly adding tool {tool_name}, will use the newest tool in function list')
                self.function_map[tool_name] = tool
        else:
            # 调用父类的处理方法
            super()._init_tool(tool)

    def _call_llm_with_retry(self, messages: List[Message], functions=None, extra_generate_cfg=None, max_retries: int = 5) -> Iterator:
        """带重试机制的LLM调用
        
        Args:
            messages: 消息列表
            functions: 函数列表
            extra_generate_cfg: 额外生成配置
            max_retries: 最大重试次数
            
        Returns:
            LLM响应流
            
        Raises:
            Exception: 重试次数耗尽后重新抛出原始异常
        """
        for attempt in range(max_retries):
            try:
                return self._call_llm(messages=messages, functions=functions, extra_generate_cfg=extra_generate_cfg)
            except Exception as e:
                # 检查是否为可重试的错误
                if self._is_retryable_error(e) and attempt < max_retries - 1:
                    delay = 2 ** attempt  # 指数退避: 1s, 2s, 4s
                    tool_logger.warning(f"LLM调用失败 (尝试 {attempt + 1}/{max_retries})，{delay}秒后重试: {str(e)}")
                    time.sleep(delay)
                    continue
                else:
                    # 不可重试的错误或已达到最大重试次数
                    if attempt > 0:
                        tool_logger.error(f"LLM调用重试失败，已达到最大重试次数 {max_retries}")
                    raise
    
    def _run(self, messages: List[Message], lang: Literal['en', 'zh', 'ja'] = 'en', **kwargs) -> Iterator[List[Message]]:

        message_list = copy.deepcopy(messages)
        response = []
        
        # 保持原有的最大调用次数限制
        total_num_llm_calls_available = self.MAX_LLM_CALL_PER_RUN if hasattr(self, 'MAX_LLM_CALL_PER_RUN') else 100
        num_llm_calls_available = total_num_llm_calls_available
        while num_llm_calls_available > 0:
            num_llm_calls_available -= 1
            extra_generate_cfg = {'lang': lang}
            if kwargs.get('seed') is not None:
                extra_generate_cfg['seed'] = kwargs['seed']
                
            output_stream = self._call_llm_with_retry(messages=message_list,
                                                functions=[func.function for func in self.function_map.values()],
                                                extra_generate_cfg=extra_generate_cfg)
            output: List[Message] = []
            for output in output_stream:
                if output:
                    yield response + output
                    
            if output:
                response.extend(output)
                message_list.extend(output)
                
                # 处理工具调用
                used_any_tool = False
                for out in output:
                    use_tool, tool_name, tool_args, _ = self._detect_tool(out)
                    if use_tool:
                        tool_result = self._call_tool(tool_name, tool_args, messages=message_list, **kwargs)

                        # 验证工具结果
                        if not tool_result:
                            tool_logger.warning(f"工具 {tool_name} 返回空结果")
                            tool_result = f"Tool {tool_name} completed execution but returned empty result"
                        elif tool_result.startswith('An error occurred when calling tool') or tool_result.startswith('工具调用失败'):
                            tool_logger.error(f"工具 {tool_name} 调用失败: {tool_result}")

                        fn_msg = Message(role=FUNCTION,
                                       name=tool_name,
                                       content=tool_result,
                                       extra={'function_id': out.extra.get('function_id', '1')})
                        message_list.append(fn_msg)
                        response.append(fn_msg)
                        yield response
                        used_any_tool = True
                
                # 如果使用了工具，继续循环
                if not used_any_tool:
                    break

        # 检查是否因为调用次数用完而退出循环
        if num_llm_calls_available == 0:
            # 根据语言选择错误消息
            if lang == 'zh':
                error_message = "工具调用超出限制"
            elif lang == 'ja':
                error_message = "ツール呼び出しが制限を超えました。"
            else:
                error_message = "Tool calls exceeded limit"
            tool_logger.error(error_message)

            error_msg = Message(
                role=ASSISTANT,
                content=error_message,
            )
            response.append(error_msg)

        yield response


# Utility functions
def read_system_prompt():
    """读取通用的无状态系统prompt"""
    with open("./prompt/system_prompt_default.md", "r", encoding="utf-8") as f:
        return f.read().strip()


def read_mcp_settings():
    """读取MCP工具配置"""
    with open("./mcp/mcp_settings.json", "r") as f:
        mcp_settings_json = json.load(f)
    return mcp_settings_json


def update_agent_llm(agent, model_name: str, api_key: str = None, model_server: str = None, generate_cfg: Dict = None):
    """动态更新助手实例的LLM和配置，支持从接口传入参数"""
    # 获取基础配置
    llm_config = {
        "model": model_name,
        "api_key": api_key,
        "model_server": model_server,
        "generate_cfg": generate_cfg if generate_cfg else {}
    }
    
    # 创建LLM实例
    llm_instance = TextChatAtOAI(llm_config)
    
    # 动态设置LLM
    agent.llm = llm_instance
    return agent


# 向后兼容：保持原有的初始化函数接口
def init_modified_agent_service_with_files(rag_cfg=None, 
                                         model_name="qwen3-next", api_key=None, 
                                         model_server=None, generate_cfg=None,
                                         system_prompt=None, mcp=None):
    """创建支持预加载文件的修改版助手实例"""
    system = system_prompt if system_prompt else read_system_prompt()
    tools = mcp if mcp else read_mcp_settings()

    llm_config = {
        "model": model_name,
        "api_key": api_key,
        "model_server": model_server,
        "generate_cfg": generate_cfg if generate_cfg else {}
    }
    
    # 创建LLM实例
    llm_instance = TextChatAtOAI(llm_config)

    bot = ModifiedAssistant(
        llm=llm_instance,
        name="修改版数据检索助手",
        description="基于智能判断循环终止的助手",
        system_message=system,
        function_list=tools,
    )

    return bot