"""Capture per-page OCR text, bboxes, and latency from the vLLM OCR endpoint."""
from __future__ import annotations

import argparse
import base64
import json
import sys
from pathlib import Path
from typing import Any, Dict

import requests


def extract_capture(body: Dict[str, Any]) -> Dict[str, Any]:
    data = body.get("data") or {}
    if not isinstance(data, dict):
        data = {}
    blocks = ((data.get("ocr_text_json") or {}).get("blocks")) or []
    boxes = [list(b["bbox"]) for b in blocks if isinstance(b, dict) and b.get("bbox")]
    return {
        "text": data.get("text_lines", "") or "",
        "boxes": boxes,
        "elapsed_seconds": data.get("elapsed_seconds"),
    }


def write_capture(cap: Dict[str, Any], out_dir: Path, stem: str) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / f"{stem}.txt").write_text(cap["text"], encoding="utf-8")
    (out_dir / f"{stem}.json").write_text(
        json.dumps({"boxes": cap["boxes"], "elapsed_seconds": cap["elapsed_seconds"]}),
        encoding="utf-8",
    )


def _payload(image_path: Path) -> Dict[str, Any]:
    suffix = image_path.suffix.lstrip(".").lower() or "png"
    return {
        "file": base64.b64encode(image_path.read_bytes()).decode("utf-8"),
        "type": "jpg" if suffix == "jpeg" else suffix,
        "skip_text_detection": False,
        "skip_table_detection": False,
        "recognize_math": False,
        "ocr_with_boxes": True,
    }


def capture_page(url: str, image_path: Path, out_dir: Path, timeout: float) -> Dict[str, Any]:
    resp = requests.post(url, json=_payload(image_path), timeout=timeout)
    cap = extract_capture(resp.json())
    if cap["elapsed_seconds"] is None:
        # Latency is a headline metric; a successful response with no timing means
        # the response contract changed. Surface it loudly instead of silently
        # dropping the page from the latency stats.
        print(
            f"WARNING: no elapsed_seconds in OCR response for {image_path.name}; "
            "latency for this page will be missing",
            file=sys.stderr,
        )
    write_capture(cap, out_dir, image_path.stem)
    return cap


def main() -> None:
    parser = argparse.ArgumentParser(description="Capture OCR text+boxes+latency per image.")
    parser.add_argument("images", nargs="+", type=Path)
    parser.add_argument("--url", required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--timeout", type=float, default=900)
    args = parser.parse_args()
    for image in args.images:
        cap = capture_page(args.url, image, args.out_dir, args.timeout)
        print(f"captured {image.name}: {len(cap['boxes'])} boxes, {cap['elapsed_seconds']}s")


if __name__ == "__main__":
    main()