import os
from typing import Dict, List, Optional
from qwen_agent.agents import Assistant
from qwen_agent.tools import BaseTool

class Qwen3Agent:
    """優化后的Qwen3智能體封裝類"""
    
    def __init__(
        self,
        model_name: str = 'Qwen3-30B-A3B',
        endpoint_type: str = 'local',  # 'dashscope' or 'local'
        custom_endpoint: Optional[str] = None,
        enable_tools: bool = True
    ):
        """
        初始化Qwen3智能體
        :param model_name: 模型名稱
        :param endpoint_type: 服務端點類型 ('dashscope' 或 'local')
        :param custom_endpoint: 自定義OpenAI兼容端點URL
        :param enable_tools: 是否啟用工具功能
        """
        self.llm_cfg = self._init_llm_config(model_name, endpoint_type, custom_endpoint)
        self.tools = self._init_tools() if enable_tools else []
        self.agent = Assistant(llm=self.llm_cfg, function_list=self.tools)
        
    def _init_llm_config(
        self,
        model_name: str,
        endpoint_type: str,
        custom_endpoint: Optional[str]
    ) -> Dict:
        """初始化LLM配置"""
        config = {
            'model': model_name,
            'generate_cfg': {
                'thought_in_content': False,  # 推薦分離思維鏈和響應內容
                'max_tokens': 2048,
                'temperature': 0.7
            }
        }
        
        if endpoint_type == 'dashscope':
            config.update({
                'model_type': 'qwen_dashscope',
                'api_key': os.getenv('DASHSCOPE_API_KEY')
            })
        else:
            config.update({
                'model_server': custom_endpoint or'http://localhost:8000/v1',
                'api_key': 'EMPTY'
            })
            
        return config
    
    def _init_tools(self) -> List[BaseTool]:
        """初始化工具配置"""
        return [
            {
                'mcpServers': {
                    'time': {
                        'command': 'uvx',
                        'args': ['mcp-server-time', '--local-timeznotallow=Asia/Shanghai']
                    },
                    "fetch": {
                        "command": "uvx",
                        "args": ["mcp-server-fetch"]
                    }
                }
            },
            'code_interpreter',  # 內置代碼解釋器
            'web_search',       # 新增網絡搜索工具
        ]
    
    def stream_response(
        self,
        messages: List[Dict],
        timeout: int = 30
    ) -> Dict:
        """
        流式響應生成
        :param messages: 對話消息列表
        :param timeout: 超時時間(秒)
        :return: 最終響應字典
        """
        try:
            for response in self.agent.run(messages=messages, timeout=timeout):
                # 可在此處添加實時處理邏輯
                pass
            return response
        except Exception as e:
            self._handle_error(e)
            return {'error': str(e)}
    
    def _handle_error(self, error: Exception):
        """統一錯誤處理"""
        # 可擴展為特定錯誤類型的處理
        print(f"Agent Error: {str(error)}")
        # 可添加日志記錄或警報邏輯

# 使用示例
if __name__ == "__main__":
    # 初始化智能體 (生產環境建議使用環境變量配置API密鑰)
    agent = Qwen3Agent(
        model_name='Qwen3-30B-A3B',
        endpoint_type='local',
        custom_endpoint='http://api.example.com/v1'
    )
    
    # 構造對話消息
    messages = [{
        'role': 'user',
        'content': 'https://qwenlm.github.io/blog/ 介紹Qwen的最新發展動態'
    }]
    
    # 獲取響應
    response = agent.stream_response(messages)
    
    # 處理響應
    if'error'notin response:
        print("完整響應:")
        print(response)
    else:
        print("請求失敗:", response['error'])