
import utils,json,requests, os
import re  # 添加正则表达式库

class FastGPTAi:

    def __init__(self, api_key: str=None, api_url: str=None):
        self._api_key = api_key or utils.get_config_value('fastgpt.api_key')
        self._api_url = api_url or utils.get_config_value('fastgpt.api_url')
        self._logger = utils.get_logger()
        self._headers ={}

    def call_ai(self, msg: str, api_key: str=None) -> json:
        try:
            api_key = api_key if api_key else self._api_key
            if api_key is None:
                self._logger.error("fastgpt.api_key 没有配置")
                raise Exception("fastgpt.api_key 没有配置")

            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            url = f"{self._api_url}/v1/chat/completions"
            data = {
                "stream": False,
                "detail": False,
                "messages": [
                    {
                        "role": "user",
                        "content": msg
                    }
                ],
                "response_format": {
                    "type": "json_object"
                }
            }

            return self.process_response(data, headers, url)
        except Exception as e:
            self._logger.error(f"Error: {str(e)}")
            return None

    def call_ai_with_image(self,image_path, msg: str,api_key: str=None) -> json:
        self._logger.info("调用 fastgpt 的AI_Image接口")
        try:
            if not os.path.exists(image_path):
                utils.get_logger().error(f"图片文件不存在: {image_path}")
                raise Exception(f"图片文件不存在: {image_path}")
            api_key = api_key if api_key else  self._api_key
            if api_key is None:
                self._logger.error("fastgpt.api_key 没有配置")
                raise Exception("fastgpt.api_key 没有配置")

            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            url = f"{self._api_url}/v1/chat/completions"
            base64_str = utils.encode_file(image_path)
            data = {
                # "model":"",
                "stream": False,
                "detail": False,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": msg
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": base64_str
                                }
                            }
                        ]
                    }
                ],
                "response_format": {
                    "type": "json_object"
                }
            }

            return self.process_response(data, headers, url)
        except Exception as e:
            self._logger.error(f"Error: {str(e)}")
            return None

    def call_ai_with_file(self,file_path, msg: str,api_key: str=None) -> json:
        self._logger.info("调用 fastgpt 的AI_File接口")
        try:
            if not os.path.exists(file_path):
                utils.get_logger().error(f"文件不存在: {file_path}")
                raise Exception(f"文件不存在: {file_path}")
            api_key = api_key if api_key else  self._api_key
            if api_key is None:
                self._logger.error("fastgpt.api_key 没有配置")
                raise Exception("fastgpt.api_key 没有配置")

            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            url = f"{self._api_url}/v1/chat/completions"
            file_name = os.path.basename(file_path)
            base64_str = utils.encode_file(file_path)
            data = {
                # "model":"",
                "stream": False,
                "detail": False,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": msg
                            },
                            {
                                "type": "file_url",
                                "name": file_name,
                                "url": base64_str
                            }
                        ]
                    }
                ]
            }

            res= self.process_response(data, headers, url)
            if res:
                return res
            return None
        except Exception as e:
            self._logger.error(f"Error: {str(e)}")
            return None

    def process_response(self, data, headers, url)->json:
        response = requests.post(url, headers=headers, json=data)
        if response.status_code == 200:
            result = response.json()
            self._logger.info(f"Response: {result}")
            content = result.get("choices", [{}])[0].get("message", {}).get("content", "")
            if content:
                if "```json" in content:
                    # 使用正则表达式提取代码块内容
                    code_block_match = re.search(r'```(json)?(.*?)```', content, re.DOTALL)
                    if code_block_match:
                        content = code_block_match.group(2).strip()
                    content = content.strip()
                try:
                    data = json.loads(content)
                    self._logger.info(f"Response_JSON: {data}")
                    return data
                except json.JSONDecodeError as e:
                    self._logger.error(f"JSON解析失败，尝试清理特殊字符: {e}")
                    # 尝试清理特殊字符后重新解析
                    content = re.sub(r'[“”‘’\x00-\x1f]', '', content)
                    try:
                        return json.loads(content)
                    except json.JSONDecodeError:
                        self._logger.error("最终JSON解析失败，返回原始内容")
                        return {'raw_content': content}
            return None
        else:
            error_msg = f"Error: {response.status_code} - {response.text}"
            self._logger.error(error_msg)
            return None
