import requests
import os
import base64
import json

# 1. 准备请求
url = "http://127.0.0.1:8000/parse/pdf" # 你的API地址
files = {'file': open('test.pdf', 'rb')}
data = {'output_format': 'json'} # 请求JSON格式以便获取图片详情

# 2. 发送请求
response = requests.post(url, files=files, data=data)
result = response.json()

# 3. 处理结果 (根据实际API返回结构调整)
# 假设返回结构中包含 markdown 内容和图片列表
md_content = result.get('markdown', '')
images = result.get('images', []) # 假设这里有图片信息

# 4. 保存图片
os.makedirs("output_images", exist_ok=True)
for img in images:
    img_name = img.get('name')
    img_data = img.get('data') # 可能是 base64 或 url
    
    # 如果是 base64
    if img_data.startswith('data:image'):
        img_bytes = base64.b64decode(img_data.split(',')[1])
        with open(f"output_images/{img_name}", 'wb') as f:
            f.write(img_bytes)
    
    # 如果是 URL (需要二次下载)
    elif img_data.startswith('http'):
        img_resp = requests.get(img_data)
        with open(f"output_images/{img_name}", 'wb') as f:
            f.write(img_resp.content)

# 5. 保存 Markdown
with open("output.md", "w", encoding="utf-8") as f:
    f.write(md_content)

print("处理完成，图片和Markdown已保存。")