#!/usr/bin/env python3
"""Generate a Doubleword batch JSONL file with deterministic Picsum image URLs."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from urllib.parse import quote

DEFAULT_MODEL = "Qwen/Qwen3-VL-235B-A22B-Instruct-FP8"
DEFAULT_PROMPT = "<|grounding|>Convert the document to markdown."
DEFAULT_OUTPUT = Path(".generated/doubleword/picsum_markdown_batch.jsonl")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--count", type=int, default=10, help="Number of requests to write.")
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help="Vision-capable model name to place in each batch body.",
    )
    parser.add_argument(
        "--prompt",
        default=DEFAULT_PROMPT,
        help="User text prompt to place before the image_url content item.",
    )
    parser.add_argument(
        "--seed-prefix",
        default="picsum-doc",
        help="Prefix used to build deterministic Picsum seeds.",
    )
    parser.add_argument(
        "--start-index",
        type=int,
        default=1,
        help="Starting integer suffix for seed generation.",
    )
    parser.add_argument("--width", type=int, default=200)
    parser.add_argument("--height", type=int, default=300)
    return parser.parse_args()


def build_seed(seed_prefix: str, index: int) -> str:
    return f"{seed_prefix}-{index:04d}"


def build_image_url(*, seed: str, width: int, height: int) -> str:
    return f"https://picsum.photos/seed/{quote(seed, safe='')}/{width}/{height}"


def build_request(*, model: str, prompt: str, image_url: str, index: int) -> dict:
    return {
        "custom_id": f"picsum-markdown-{index:04d}",
        "method": "POST",
        "url": "/v1/chat/completions",
        "body": {
            "model": model,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {"type": "image_url", "image_url": {"url": image_url}},
                    ],
                }
            ],
        },
    }


def main() -> None:
    args = parse_args()
    if args.count <= 0:
        raise SystemExit("--count must be greater than 0.")

    args.output.parent.mkdir(parents=True, exist_ok=True)

    with args.output.open("w", encoding="utf-8") as handle:
        for offset in range(args.count):
            index = args.start_index + offset
            seed = build_seed(args.seed_prefix, index)
            image_url = build_image_url(seed=seed, width=args.width, height=args.height)
            request = build_request(
                model=args.model,
                prompt=args.prompt,
                image_url=image_url,
                index=index,
            )
            handle.write(json.dumps(request, separators=(",", ":")))
            handle.write("\n")

    print(f"Wrote {args.count} requests to {args.output}")


if __name__ == "__main__":
    main()
