#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "typer>=0.12",
#     "loguru>=0.7.2",
#     "transformers>=5.15",
#     "torch>=2.10",
#     "torchvision>=0.25",
#     "accelerate>=1.0",
#     "peft>=0.19.1",
#     "pillow>=12.2",
# ]
# ///
"""sVLM-Council — run the same image+prompt through small local VLMs.

All members run locally via transformers (no Spaces, no Inference API) through
one generic `query()` function; per-model quirks live in the COUNCIL registry.
Works for data extraction or vanilla VQA alike.

Council members (all in-tree transformers, dense, MPS-friendly):
- qianfan-ocr     -> baidu/Qianfan-OCR                  (4.7B, KIE leader)
- hunyuan-ocr     -> tencent/HunyuanOCR                 (1.1B, scene text + IE)
- granite-vision  -> ibm-granite/granite-vision-4.1-4b  (4B, key-value extraction)
- qwen3-vl-8b     -> Qwen/Qwen3-VL-8B-Instruct          (8.8B, best sub-10B OCRBench)

Models are loaded one at a time (single-slot cache with eviction): the four
together exceed 32GB unified memory, and first use downloads each checkpoint.

Usage:
    uv run svlm_council.py bib.jpg "What is the bib number in this image?"
    uv run svlm_council.py bib.jpg "Describe this image." -m hunyuan-ocr -t 0.7
    uv run svlm_council.py bib.jpg "..." -s "You are a terse assistant." --json -o out.json

As a library (e.g. from extract_arena.py):
    import svlm_council
    text = svlm_council.query("qianfan-ocr", "bib.jpg", "What is the bib number?")

Results go to stdout; logs go to stderr (loguru default).
"""
from __future__ import annotations

import gc
import json
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Optional

import typer
from loguru import logger

# ----------------------------------------------------------------- registry


@dataclass(frozen=True)
class CouncilSpec:
    name: str
    model_id: str
    dtype: str = "auto"  # "auto" or a torch dtype name like "bfloat16"
    device_map: str = "mps"  # accelerate's "auto" offloads to disk far too eagerly on Macs
    attn_implementation: str = "sdpa"  # none of the members needs flash-attn
    processor_kwargs: dict = field(default_factory=dict)
    generate_kwargs: dict = field(default_factory=dict)


COUNCIL: dict[str, CouncilSpec] = {
    "qianfan-ocr": CouncilSpec(
        "qianfan-ocr",
        "baidu/Qianfan-OCR",
        dtype="bfloat16",
    ),
    "hunyuan-ocr": CouncilSpec(
        "hunyuan-ocr",
        "tencent/HunyuanOCR",
        dtype="bfloat16",
        attn_implementation="eager",  # recommended by the transformers doc for the OCR path
        processor_kwargs={"backend": "pil"},
        generate_kwargs={"repetition_penalty": 1.08},  # model card's recommended setting
    ),
    "granite-vision": CouncilSpec(
        "granite-vision",
        "ibm-granite/granite-vision-4.1-4b",
        dtype="bfloat16",
    ),
    "qwen3-vl-8b": CouncilSpec(
        "qwen3-vl-8b",
        "Qwen/Qwen3-VL-8B-Instruct",
    ),
}


def _resolve_spec(model: str) -> CouncilSpec:
    # Any HF repo id also works: unknown names get default handling, so trying
    # a new model is just a name change.
    return COUNCIL.get(model) or CouncilSpec(name=model, model_id=model)


# ------------------------------------------------- single-slot model cache
# The council members total ~37GB bf16 — more than this host's unified memory —
# so only one model stays resident; switching models evicts the previous one.

_LOADED: dict[str, object] = {}


def _load(spec: CouncilSpec):
    import torch
    from transformers import AutoModelForImageTextToText, AutoProcessor

    if _LOADED.get("name") == spec.name:
        return _LOADED["processor"], _LOADED["model"]
    if _LOADED:
        logger.info("Evicting {} to free memory", _LOADED["name"])
        _LOADED.clear()
        gc.collect()
        if torch.backends.mps.is_available():
            torch.mps.empty_cache()

    logger.info("Loading {} ({}) — first use downloads the weights", spec.name, spec.model_id)
    processor = AutoProcessor.from_pretrained(spec.model_id, **spec.processor_kwargs)
    device_map = spec.device_map if torch.backends.mps.is_available() else "auto"
    model = AutoModelForImageTextToText.from_pretrained(
        spec.model_id,
        dtype=spec.dtype if spec.dtype == "auto" else getattr(torch, spec.dtype),
        device_map=device_map,
        attn_implementation=spec.attn_implementation,
    ).eval()
    _LOADED.update(name=spec.name, processor=processor, model=model)
    return processor, model


# ----------------------------------------------------------------- core API


def query(
    model: str,
    image_path: Path | str,
    prompt: str,
    *,
    system_prompt: str | None = None,
    temperature: float = 0.0,
    max_tokens: int = 1024,
) -> str:
    """Ask one council member (or any HF repo id) a question about an image.

    Generic over tasks: works for targeted data extraction and vanilla VQA.
    Returns the model's text reply; raises on failure (callers isolate errors).
    """
    import torch

    spec = _resolve_spec(model)
    processor, vlm = _load(spec)

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
    messages.append(
        {
            "role": "user",
            "content": [
                {"type": "image", "url": str(image_path)},
                {"type": "text", "text": prompt},
            ],
        }
    )
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(vlm.device)
    inputs.pop("token_type_ids", None)  # required for Qwen3-VL, harmless for the rest

    gen_kwargs: dict = {
        "max_new_tokens": max_tokens,
        "use_cache": True,  # Qianfan-OCR's config ships use_cache: false
        **spec.generate_kwargs,
    }
    if temperature > 0:
        gen_kwargs.update(do_sample=True, temperature=temperature)
    else:
        # Explicit Nones neutralize sampling defaults baked into some
        # generation_configs (Qwen3-VL ships do_sample=true, temperature=0.7).
        gen_kwargs.update(do_sample=False, temperature=None, top_p=None, top_k=None)

    with torch.inference_mode():
        generated = vlm.generate(**inputs, **gen_kwargs)
    return processor.batch_decode(
        generated[:, inputs["input_ids"].shape[1] :],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,  # preserve whitespace fidelity in OCR-ish output
    )[0].strip()


# ------------------------------------------------------------- orchestration


@dataclass
class CouncilResult:
    model: str
    text: str = ""
    latency_s: float = 0.0
    error: str | None = None


def run_council(
    image_path: Path,
    prompt: str,
    model_names: list[str],
    *,
    system_prompt: str | None = None,
    temperature: float = 0.0,
    max_tokens: int = 1024,
) -> list[CouncilResult]:
    results = []
    for name in model_names:
        logger.info("[{}] prompt: {!r}", name, prompt)
        t0 = time.perf_counter()
        try:
            text = query(
                name,
                image_path,
                prompt,
                system_prompt=system_prompt,
                temperature=temperature,
                max_tokens=max_tokens,
            )
            result = CouncilResult(name, text=text, latency_s=round(time.perf_counter() - t0, 2))
            logger.info("[{}] done in {}s", name, result.latency_s)
        except Exception as exc:  # noqa: BLE001 — isolate failures per model
            logger.exception("[{}] failed", name)
            result = CouncilResult(
                name,
                error=f"{type(exc).__name__}: {exc}",
                latency_s=round(time.perf_counter() - t0, 2),
            )
        results.append(result)
    return results


# -------------------------------------------------------------- presentation


def results_to_payload(
    image_path: Path,
    prompt: str,
    results: list[CouncilResult],
    *,
    system_prompt: str | None = None,
    temperature: float = 0.0,
    max_tokens: int = 1024,
) -> dict:
    return {
        "image": str(image_path),
        "prompt": prompt,
        "system_prompt": system_prompt,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "results": [asdict(r) for r in results],
        "errors": [f"{r.model}: {r.error}" for r in results if r.error],
    }


def print_results(results: list[CouncilResult]) -> None:
    for r in results:
        status = "── ERROR " if r.error else ""
        print("\n" + f"── {r.model} {status}── {r.latency_s}s ".ljust(60, "─"))
        print(r.error if r.error else r.text)


# ---------------------------------------------------------------------- CLI

cli = typer.Typer(add_completion=False, no_args_is_help=True)


@cli.command()
def main(
    image: Path = typer.Argument(..., exists=True, dir_okay=False, readable=True, help="Input image"),
    prompt: str = typer.Argument(..., help="Question or instruction for the image"),
    models: Optional[list[str]] = typer.Option(
        None, "--models", "-m", help="Models to run (repeatable): council names or HF repo ids. Council: " + ", ".join(COUNCIL)
    ),
    system_prompt: Optional[str] = typer.Option(None, "--system-prompt", "-s", help="Optional system prompt"),
    temperature: float = typer.Option(0.0, "--temperature", "-t", min=0.0, help="0 = deterministic (greedy)"),
    max_tokens: int = typer.Option(1024, "--max-tokens", min=1, help="Max new tokens to generate"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Also write results as JSON to this file"),
    as_json: bool = typer.Option(False, "--json", help="Print JSON to stdout instead of readable text"),
) -> None:
    """Ask every selected local VLM the same question about the image and compare answers."""
    names = models or list(COUNCIL)
    results = run_council(
        image, prompt, names, system_prompt=system_prompt, temperature=temperature, max_tokens=max_tokens
    )
    payload = results_to_payload(
        image, prompt, results, system_prompt=system_prompt, temperature=temperature, max_tokens=max_tokens
    )

    if as_json:
        print(json.dumps(payload, indent=2, ensure_ascii=False))
    else:
        print_results(results)
    if output:
        output.write_text(json.dumps(payload, indent=2, ensure_ascii=False))
        logger.info("Wrote {}", output)

    if all(r.error for r in results):
        raise typer.Exit(1)


if __name__ == "__main__":
    cli()