import dashscope


class QwenClient:
    def __init__(self):
        self.api_key = "sk-fef6289a33024fcca98edf6fae3afbcc"

    def chat(self, model="qwen3-max", system_prompt="You are a helpful assistant.", user_prompt=""):
        """
        普通聊天，不使用搜索功能

        Args:
            model: 模型名称，默认为qwen3-max
            system_prompt: 系统提示词
            user_prompt: 用户提示词

        Returns:
            str: AI回复内容
        """
        try:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ]

            response = dashscope.Generation.call(
                api_key=self.api_key,
                model=model,
                messages=messages,
                result_format="message"
            )

            if response.status_code != 200:
                raise Exception(f"API调用失败: {response.message}")

            return response["output"]["choices"][0]["message"]["content"]

        except Exception as e:
            raise Exception(f"QwenClient chat失败: {str(e)}")

    def search_and_chat(self, model="qwen3-max", system_prompt="You are a helpful assistant.", user_prompt="", search_strategy="max"):
        """
        搜索并聊天

        Args:
            model: 模型名称，默认为qwen3-max
            system_prompt: 系统提示词
            user_prompt: 用户提示词
            search_strategy: 搜索策略，可选值: turbo, max, agent

        Returns:
            dict: 包含回复内容和搜索结果
        """
        try:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ]

            response = dashscope.Generation.call(
                api_key=self.api_key,
                model=model,
                messages=messages,
                enable_search=True,
                search_options={
                    "forced_search": True,
                    "enable_source": True,
                    "search_strategy": search_strategy
                },
                result_format="message"
            )

            if response.status_code != 200:
                raise Exception(f"API调用失败: {response.message}")

            content = response["output"]["choices"][0]["message"]["content"]
            search_results = []

            if hasattr(response.output, 'search_info') and response.output.search_info:
                search_results = response.output.search_info.get("search_results", [])

            return {
                "content": content,
                "search_results": search_results
            }

        except Exception as e:
            raise Exception(f"QwenClient search_and_chat失败: {str(e)}")


if __name__ == "__main__":
    client = QwenClient()

    # 测试
    try:

#         prompt = """请将工具 [小红书] 的 [帖子详情] 功能翻译为英文，并返回翻译后的英文名称
# 要求：
# 1. 要将工具和功能一起翻译，如果没有合适的翻译，用中文拼音替代
# 2. 翻译返回的英文名称所有单词用下划线拼接，不能出现空格。比如: chatgpt_ai_search
# 3. 仅输出以下 JSON，不添加任何其他文字、注释或解释：
# {
#     "english_name": ""
# }"""
#         result = client.chat(user_prompt=prompt)
#         print(result)

        user_prompt = """明星宠物 造型元素 规律"""

        
        # user_prompt = "请搜索 白瓜AI 官网"
        
        result = client.search_and_chat(user_prompt=user_prompt, search_strategy="agent")

        print("="*20 + "搜索结果" + "="*20)
        for web in result["search_results"]:
            print(f"[{web['index']}]: [{web['title']}]({web['url']})")

        print("="*20 + "回复内容" + "="*20)
        print(result["content"])

    except Exception as e:
        print(f"错误: {e}")