import base64
import io
import sys

import requests
from PIL import Image

API_URL = "https://api.1808366.xyz/v1/chat/completions"
MODEL = "qwen3.5-9b-uncensored-nothink"


QUALITY = 85          # 高画质，保留绝大部分细节
MAX_DIM = 1280        # 最长边上限，保证细节清晰的最小合理尺寸（一次缩放，不迭代）


def img_to_data_url(path):
    img = Image.open(path).convert("RGB")
    # 仅在超长边过大时一次性缩到目标尺寸，保证细节清晰的同时控制体积
    if max(img.size) > MAX_DIM:
        ratio = MAX_DIM / max(img.size)
        img = img.resize(
            (int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS
        )
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=QUALITY)
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/jpeg;base64,{b64}"


def ask(path, question):
    payload = {
        "model": MODEL,
        "messages": [{
            "role": "user",
            "content": [
                {"type": "text", "text": question},
                {"type": "image_url", "image_url": {"url": img_to_data_url(path)}}
            ]
        }]
    }
    r = requests.post(API_URL, json=payload, timeout=120)
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"]


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("用法: python scripts/img_query.py <图片路径> <问题> [输出文件]")
        sys.exit(1)
    path, question = sys.argv[1], sys.argv[2]
    out = sys.argv[3] if len(sys.argv) > 3 else None
    result = ask(path, question)
    if out:
        with open(out, "w", encoding="utf-8") as f:
            f.write(result)
    else:
        sys.stdout.write(result)