# -*- coding: utf-8 -*-
"""
Created on 2024-10-31 
---------
@summary:  智普文件抽取服务
---------
@author: Dzr
"""
import base64
import json
import pathlib
from pathlib import Path

from zhipuai import ZhipuAI


class ZhipuFileExtract:

    def __init__(self, api_key):
        self._client = ZhipuAI(
            api_key=api_key,
            base_url="https://open.bigmodel.cn/api/paas/v4"
        )

    def file_extract(self, file):
        file_object = self.create(file)
        file_content = json.loads(self.get_content(file_object.id))["content"]
        self.delete(file_object.id)
        print('智普ocr > ', file_content)
        return file_content

    def create(self, file):
        # 格式限制：.PDF .DOCX .DOC .XLS .XLSX .PPT .PPTX .PNG .JPG .JPEG .CSV .PY .TXT .MD .BMP .GIF
        # 大小：单个文件50M、总数限制为100个文件
        # 支持retrieval、batch、fine-tune、file-extract文件
        return self._client.files.create(file=file, purpose="file-extract")

    def get_content(self, file_id):
        return self._client.files.content(file_id=file_id).content

    def delete(self, file_id):
        self._client.files.delete(file_id=file_id)

    def glm_4v_flash(self, file, text=None):
        if isinstance(file, pathlib.Path):
            img_bytes = file.read_bytes()
            img_name = file.name
        else:
            img_name, img_bytes, _ = file

        if text is None:
            text = "这是一张字体图片,图片内容是什么？你只需要返回图中内容即可,不要解释,不要返回任何其它内容"

        img_base = base64.b64encode(img_bytes).decode('utf-8')
        response = self._client.chat.completions.create(
            model="glm-4v-flash",  # 填写需要调用的模型名称
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": img_base
                            }
                        },
                        {
                            "type": "text",
                            "text": text
                        }
                    ]
                }
            ]
        )
        print(img_name, "--> ", response.choices[0].message)
        return response.choices[0].message.content


if __name__ == '__main__':
    client = ZhipuFileExtract(
        api_key=""
    )
    ret = client.glm_4v_flash(Path('download.png'), text='这是一张100以内数学计算图片验证码，图片内容计算结果是什么？你只需要返回图中内容计算结果即可,不要解释,不要返回任何其它内容')
    print(ret)