import base64

from openai import OpenAI


# 读取本地图片文件并转换为Base64编码
def encode_image_to_base64(image_path, ):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


def getAiOcr(prompt, image_path):
    # 初始化OpenAI客户端
    client = OpenAI(
        # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
        api_key="sk-383770c29f4c41f0bec52a640af80e64",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )

    # 本地图片路径
    if prompt is None:
        image_path = "G:\\wd\\sfz33.jpg"  # 替换为你的本地图片路径

    # 将本地图片编码为Base64
    image_base64 = encode_image_to_base64(image_path)

    # 创建请求
    if prompt is None:
        prompt = "Convert the provided image of a PDF document strictly into valid Markdown. If there are mathematical formulas in the document, use Mathjax."

    completion = client.chat.completions.create(
        model="qwen-vl-plus",  # 此处以qwen-vl-plus为例，可按需更换模型名称。
        messages=[
            {"role": "system", "content": "You are a tool to parse documents."},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}  # 使用Base64编码的图片数据
                    }
                ]
            }
        ]
    )

    # 打印结果
    print(completion.model_dump_json())
    print(completion.choices[0].message.content)
    return completion.choices[0].message.content


def getAiOcrBase64(prompt, image_base64):
    # 初始化OpenAI客户端
    client = OpenAI(
        # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
        api_key="sk-383770c29f4c41f0bec52a640af80e64",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )

    # 本地图片路径
    if prompt is None:
        image_path = "G:\\wd\\sfz33.jpg"  # 替换为你的本地图片路径

    # 将本地图片编码为Base64
    # image_base64 = encode_image_to_base64(image_path)

    # 创建请求
    if prompt is None:
        prompt = "Convert the provided image of a PDF document strictly into valid Markdown. If there are mathematical formulas in the document, use Mathjax."

    completion = client.chat.completions.create(
        model="qwen-vl-plus",  # 此处以qwen-vl-plus为例，可按需更换模型名称。
        messages=[
            {"role": "system", "content": "You are a tool to parse documents."},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}  # 使用Base64编码的图片数据
                    }
                ]
            }
        ]
    )

    # 打印结果

    return completion.choices[0].message.content


if __name__ == "__main__":
    getAiOcr('提取有效信息', "G:\\wd\\sfz33.jpg")
