"""Read page images through `serve.sh`, one request per page.

    python examples/client.py page1.png page2.jpg ...                    # Markdown (transcription mode)
    python examples/client.py --mode grounding page1.png                  # Markdown with block boxes
    python examples/client.py --url http://127.0.0.1:8000/v1 --out pages.jsonl --concurrency 64 *.png

Pages are downscaled to a 2,048 px longest edge (the training and evaluation resolution; 1,540 px is
faster and slightly less accurate). Greedy decoding, at most 4,096 output tokens. A page that starts to
loop is trimmed after the first copy of the repeated span (loop_guard.scan), finish = "loop".
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import io
import json
import os
import sys

from openai import AsyncOpenAI
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from loop_guard import scan  # noqa: E402


def encode(path: str, edge: int) -> str:
    im = Image.open(path).convert("RGB")
    if max(im.size) > edge:
        im.thumbnail((edge, edge), Image.LANCZOS)
    buf = io.BytesIO()
    im.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


async def read(cli: AsyncOpenAI, path: str, mode: str, edge: int, sem: asyncio.Semaphore) -> dict:
    async with sem:
        b64 = await asyncio.to_thread(encode, path, edge)
        content = [{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}]
        if mode == "grounding":
            content.append({"type": "text", "text": "grounding"})
        resp = await cli.chat.completions.create(model="peppercorn", temperature=0, max_tokens=4096,
                                                 messages=[{"role": "user", "content": content}])
    text, finish = resp.choices[0].message.content or "", resp.choices[0].finish_reason
    cut = scan(text)
    if cut is not None:
        text, finish = text[:cut], "loop"
    return {"image": path, "text": text, "finish": finish, "tokens": resp.usage.completion_tokens}


async def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("images", nargs="+")
    ap.add_argument("--url", default="http://127.0.0.1:8000/v1")
    ap.add_argument("--mode", choices=["text", "grounding"], default="text")
    ap.add_argument("--edge", type=int, default=2048)
    ap.add_argument("--concurrency", type=int, default=32)
    ap.add_argument("--out", help="write one JSON line per page here instead of printing Markdown")
    a = ap.parse_args()
    cli = AsyncOpenAI(base_url=a.url, api_key="x", timeout=3600)
    sem = asyncio.Semaphore(a.concurrency)
    rows = await asyncio.gather(*(read(cli, p, a.mode, a.edge, sem) for p in a.images))
    if a.out:
        with open(a.out, "w", encoding="utf-8") as fh:
            for r in rows:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    else:
        for r in rows:
            print(f"<!-- {r['image']} finish={r['finish']} tokens={r['tokens']} -->\n{r['text']}\n")


if __name__ == "__main__":
    asyncio.run(main())