import argparse
import base64
import json
import re
from pathlib import Path

from PIL import Image


def image_block(spec, protocol):
    image_id, separator, filename = spec.partition("=")
    if not separator or not re.fullmatch(r"[A-Za-z][A-Za-z0-9_-]{0,39}", image_id):
        raise ValueError("Each image must be image_id=filename, using a short safe ID")
    path = Path(filename)
    if path.stat().st_size > 20 * 1024 * 1024:
        raise ValueError("Local example limit: 20 MiB per image")
    with Image.open(path) as image:
        if image.format not in {"PNG", "JPEG", "WEBP"} or getattr(image, "is_animated", False):
            raise ValueError("Local example accepts static PNG, JPEG or WebP only")
        if image.width * image.height > 16_000_000:
            raise ValueError("Local example limit: 16 million pixels per image")
        mime = {"PNG": "image/png", "JPEG": "image/jpeg", "WEBP": "image/webp"}[image.format]
        image.load()
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    if protocol == "anthropic":
        return [{"type": "text", "text": f"image_id: {image_id}"}, {"type": "image", "source": {"type": "base64", "media_type": mime, "data": encoded}}]
    if protocol == "responses":
        return [{"type": "input_text", "text": f"image_id: {image_id}"}, {"type": "input_image", "image_url": f"data:{mime};base64,{encoded}"}]
    return [{"type": "text", "text": f"image_id: {image_id}"}, {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{encoded}"}}]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--protocol", choices=["chat", "responses", "anthropic"], default="chat")
    parser.add_argument("--model", required=True)
    parser.add_argument("--images", nargs="+", required=True)
    parser.add_argument("--prompt", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if args.output.exists():
        raise ValueError("Output already exists; choose a new path")
    if len(args.images) > 8:
        raise ValueError("Local example limit: eight images; model limits must be checked separately")
    ids = [spec.partition("=")[0] for spec in args.images]
    if len(ids) != len(set(ids)):
        raise ValueError("Duplicate image_id")
    if args.prompt.stat().st_size > 32 * 1024:
        raise ValueError("Local example prompt limit: 32 KiB")
    prompt = args.prompt.read_text(encoding="utf-8").strip()
    if not prompt:
        raise ValueError("Prompt is empty")
    content = []
    for spec in args.images:
        content.extend(image_block(spec, args.protocol))
    content.append({"type": "input_text" if args.protocol == "responses" else "text", "text": prompt})
    body = {"model": args.model, "input" if args.protocol == "responses" else "messages": [{"role": "user", "content": content}]}
    if args.protocol == "anthropic":
        body["max_tokens"] = 1500
    with args.output.open("x", encoding="utf-8") as output:
        json.dump(body, output, ensure_ascii=False, indent=2)
        output.write("\n")
    print(json.dumps({"protocol": args.protocol, "image_ids": ids, "output": str(args.output), "network_requests": 0}, ensure_ascii=False))


if __name__ == "__main__":
    main()