#!/home/devcontainers/tinyllm/.venv/bin/python
"""DuoVLM-40M 命令行：给一张图 + 一个问题 → 回答。

  python ask_cli.py --image photo.jpg --question "How many people are in the image?"
  python ask_cli.py --image photo.jpg                      # 不给问题 = 描述这张图
  python ask_cli.py --image photo.jpg --blind              # 并排跑一遍"不看图"对照
  python ask_cli.py --dir ./photos --question "What is happening?"
  python ask_cli.py --text "The history of the city begins in" --continue-n 60

提示：40M 小模型擅长物体/场景识别、计数、是-否、简短描述；不擅长精细颜色、
读文字(OCR)、长句推理；长输出必须带 no_repeat_ngram=3（默认已开）。
"""
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from duovlm_infer import DESC_Q, DuoVLMInfer  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", action="append", default=[])
    ap.add_argument("--dir", default="")
    ap.add_argument("--question", default="")
    ap.add_argument("--questions", default="", help="同一张图连问多个，用 | 分隔")
    ap.add_argument("--text", default="", help="纯文本续写种子（不用图）")
    ap.add_argument("--continue-n", type=int, default=40)
    ap.add_argument("--blind", action="store_true")
    ap.add_argument("--no-image", action="store_true")
    ap.add_argument("--max-new", type=int, default=12)
    ap.add_argument("--ngram", type=int, default=3)
    ap.add_argument("--rep-penalty", type=float, default=1.0)
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--top-k", type=int, default=0)
    ap.add_argument("--no-loop-break", action="store_true")
    ap.add_argument("--weights", default="")
    ap.add_argument("--json", action="store_true", help="输出 JSON（给脚本/服务用）")
    a = ap.parse_args()

    imgs = list(a.image)
    if a.dir:
        imgs += [str(p) for p in sorted(Path(a.dir).glob("*.jpg"))[:20]]
    if not imgs and not a.text and not a.no_image:
        print("需要 --image / --dir（或 --text / --no-image 走无图模式）", file=sys.stderr)
        raise SystemExit(2)

    qs = [x.strip() for x in a.questions.split("|") if x.strip()] or [a.question]
    kw = dict(ngram=a.ngram, rep_penalty=a.rep_penalty, temperature=a.temperature,
              top_k=a.top_k, loop_break=not a.no_loop_break)

    vlm = DuoVLMInfer(weights=a.weights or None, verbose=not a.json)
    if not a.json:
        print(vlm.param_report())
        print("-" * 88, flush=True)
    rows = []
    if a.text:
        for seed in [x for x in a.text.split("|") if x]:
            r = vlm.continue_text(seed, max_new=a.continue_n, **kw)
            r["seed"] = seed
            rows.append(r)
            if not a.json:
                print(f"种子: {seed}\n续写: {r['answer']}\n      ({r['ms']}ms)\n", flush=True)
    else:
        for p in imgs or [None]:
            for i, q in enumerate(qs):
                r = vlm.ask(p, q, no_image=a.no_image, max_new=a.max_new, **kw)
                r.update({"image": p, "question": q or DESC_Q})
                if a.blind and p:
                    b = vlm.ask(p, q, blind=True, max_new=a.max_new, **kw)
                    r["blind_answer"] = b["answer"]
                rows.append(r)
                if not a.json:
                    tag = "" if i else Path(p).name if p else "(无图)"
                    print(f"{tag:30s} Q: {r['question']}")
                    print(f"{'':30s} [{r['source']}] {r['ms']}ms  答: {r['answer']}")
                    if "blind_answer" in r:
                        print(f"{'':30s} 不看图: {r['blind_answer']}")
    if a.json:
        print(json.dumps(rows, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()