import requests
import json


class GLMsAPIWrapper:
    def __init__(self, api_key, api_secret, assistant_id):
        self.api_key = api_key
        self.api_secret = api_secret
        self.assistant_id = assistant_id
        self.access_token = self.get_access_token(api_key, api_secret)
        self.conversation_id = None  # 用于存储会话ID

    def get_access_token(self, api_key, api_secret):
        url = "https://chatglm.cn/chatglm/assistant-api/v1/get_token"
        data = {
            "api_key": api_key,
            "api_secret": api_secret
        }
        response = requests.post(url, json=data)
        token_info = response.json()
        return token_info['result']['access_token']

    def handle_response(self, data_dict):
        message = data_dict.get("message")
        if message:
            content = message.get("content")
            if content:
                response_type = content.get("type")
                if response_type == "text":
                    return content.get("text", "No text provided")
                elif response_type == "image":
                    images = content.get("image", [])
                    image_urls = ", ".join(image.get("image_url") for image in images)
                    return image_urls
                elif response_type == "code":
                    return content.get('code')
                elif response_type == "execution_output":
                    return content.get('content')
                elif response_type == "system_error":
                    return content.get('content')
                elif response_type == "tool_calls":
                    # 检查 'tool_calls' 是否存在
                    if 'tool_calls' in data_dict:
                        return data_dict['tool_calls']
                    else:
                        return "No tool calls found in the response"
                elif response_type == "browser_result":
                    content = json.loads(content.get("content", "{}"))
                    return f"Browser Result - Title: {content.get('title')} URL: {content.get('url')}"
        # 如果没有匹配的响应类型，返回一个默认的消息
        return "Unsupported response type"

    def send_message(self, prompt, new_conversation=False, file_list=None, meta_data=None):
        url = "https://chatglm.cn/chatglm/assistant-api/v1/stream"
        headers = {
            "Authorization": f"Bearer {self.access_token}",
            "Content-Type": "application/json"
        }
        data = {
            "assistant_id": self.assistant_id,
            "prompt": prompt,
        }

        # 根据new_conversation的值决定是否使用现有的conversation_id
        if not new_conversation and self.conversation_id:
            data["conversation_id"] = self.conversation_id
        elif new_conversation:
            self.conversation_id = None  # 如果是新的对话，先清除之前的会话ID

        if file_list:
            data["file_list"] = file_list
        if meta_data:
            data["meta_data"] = meta_data

        output = None  # Initialize output variable to store the final output
        with requests.post(url, json=data, headers=headers, stream=True) as response:
            if response.status_code == 200:
                for line in response.iter_lines():
                    if line:
                        decoded_line = line.decode('utf-8')
                        if decoded_line.startswith('data:'):
                                                    decoded_line = line.decode('utf-8')
                        if decoded_line.startswith('data:'):
                            data_dict = json.loads(decoded_line[5:])
                            output = self.handle_response(data_dict)
                            # 如果是新的对话，保存返回的conversation_id
                            if 'conversation_id' in data_dict and new_conversation:
                                self.conversation_id = data_dict['conversation_id']
            else:
                output = "Request failed", response.status_code

        return {"output":output,"conversation_id":self.conversation_id}  # Return the final output after processing all lines