import requests
from io import BytesIO
import base64
import json
from PIL import Image, ImageFilter, ImageEnhance
from config import PROMPT_EXTRACT_ICON, PROMPT_EXTRACT_NAME, PROMPT_EXTRACT_COMPONENTS, PROMPT_EXTRACT_KEYWORD, PROMPT_EXTRACT_PREVENTION, PROMPT_EXTRACT_SUPPLIER

def image_to_base64(pil_image, image_format="JPEG"):
    """将PIL Image图像转换为Base64编码"""
    buffered = BytesIO()
    pil_image.save(buffered, format=image_format)
    img_byte_array = buffered.getvalue()
    encode_image = base64.b64encode(img_byte_array).decode('utf-8')
    return encode_image

def resize_image(image, max_size=512):
    """缩放图像尺寸，保持 OCR 质量"""
    width, height = image.size
    max_dim = max(width, height)

    # 如果图像不需要缩小，直接返回
    if max_dim <= max_size:
        return image

    scaling_factor = max_size / max_dim
    new_width = int(width * scaling_factor)
    new_height = int(height * scaling_factor)

    resized = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    resized = resized.filter(ImageFilter.UnsharpMask(radius=1, percent=120, threshold=3))

    enhancer = ImageEnhance.Contrast(resized)
    resized = enhancer.enhance(1.1)

    return resized

image = Image.open('./test1.jpg')
image = resize_image(image)
image_base64 = image_to_base64(image)


# response = requests.post(
#     "http://127.0.0.1:8000/api/v1/ocr",
#     json={
#         "image": image_base64,
#         "text": [PROMPT_EXTRACT_ICON, PROMPT_EXTRACT_NAME, PROMPT_EXTRACT_COMPONENTS, PROMPT_EXTRACT_KEYWORD, PROMPT_EXTRACT_PREVENTION, PROMPT_EXTRACT_SUPPLIER]
#     }
# )

response = requests.post(
    "https://u475436-9425-5ad0e9a4.gda1.seetacloud.com:6443/api/v1/agent_ocr",
    json={
        "image": image_base64,
    }
)

result = response.json()
print(result['data'])