import logging
import time
from zhipuai import ZhipuAI

logger = logging.getLogger(__name__)

class ZhipuAlltoolService:
    def __init__(self):
        self.model_name = "glm-4-alltools"
        self.app_secret_key = "d54f764a1d67c17d857bd3983b772016.GRjowY0fyiMNurLc"
        logger.info("ZhipuAlltoolService initialized with model: %s", self.model_name)

    def func_call_classify(self, message, categories:list=None):
        logger.info("Starting func_call_classify call")
        start_time = time.time()
        default_categories = ["web_search", "retrive_knowledge", "generate_report", "update_report", "clear_report"]
        categories = categories if categories else default_categories
        client = ZhipuAI(api_key=self.app_secret_key)
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "classify_user_input",
                    "description": "根据用户输入，判断用户意图，返回意图类型",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "category": {
                                "type": "string",
                                # "description": "用户的意图有以下选项:web_search,retrive_knowledge,generate_report,update_report,clear_report",
                                "description": "用户的意图有以下选项:" + ",".join(categories)
                            }
                        },
                        "required": ["category"],
                    },
                }
            }
        ]
        messages = [
            {
                "role": "user",
                "content": f"判断以下用户输入的意图并将其分类返回意图类型：{message}"
            }   
        ]
        try:
            response = client.chat.completions.create(
                model="glm-4-flash",  # 填写需要调用的模型名称
                messages= messages,
                tools= tools,
                tool_choice="auto"
            )
            print(response)
            return response.choices[0].message.tool_calls[0].function.arguments
        except Exception as e:
            logger.error("Error in web_search call: %s", str(e))
            raise e
    
    def func_call_yes_or_no(self, message, question):
        logger.info("Starting func_call_yes_or_no call")

        client = ZhipuAI(api_key=self.app_secret_key)
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "classify_user_input",
                    "description": "根据用户输入的信息和问题，回答yes或者no",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "answer": {
                                "type": "string",
                                "description": "判断结果有以下选项:yes,no"
                            }
                        },
                        "required": ["answer"],
                    },
                }
            }
        ]
        messages = [
            {
                "role": "user",
                "content": f"{question}：{message}"
            }   
        ]
        try:
            response = client.chat.completions.create(
                model="glm-4-flash",  # 填写需要调用的模型名称
                messages= messages,
                tools= tools,
                tool_choice="auto"
            )
            print(response)
            return response.choices[0].message.tool_calls[0].function.arguments
        except Exception as e:
            logger.error("Error in web_search call: %s", str(e))
            raise e


    def web_search_sse(self, message):
        logger.info("Starting web_search_sse call")
        start_time = time.time()
        client = ZhipuAI(api_key=self.app_secret_key)
        try:
            response = client.chat.completions.create(
                model="glm-4-alltools",  # 填写需要调用的模型名称
                messages=[
                    {
                        "role": "user",
                        "content":[
                            {
                                "type":"text",
                                "text":message
                            }
                        ]
                    }
                ],
                stream=True,
                tools=[
                    {
                        "type": "web_browser"
                    }
                ]
            )
            for chunk in response:
                # print(chunk)
                print("content: ",chunk.choices[0].delta.content)
                print("tool_calls: ",chunk.choices[0].delta.tool_calls)
                logger.info("content: %s", str(chunk))
                yield chunk.choices[0].delta.content
        except Exception as e:
            logger.error("Error in web_search_sse call: %s", str(e))
            raise e