#!/usr/bin/env python3
"""Tier-2 vision captioner backed by llama.cpp (GGUF) instead of PyTorch/transformers.

Drop-in replacement for tools/vision_caption.py: SAME manifest/output contract, so the Java
HighlightVisionDirector is unchanged — only the configured script path differs. Uses a stronger local VLM
(e.g. Qwen2.5-VL-3B-Instruct) through llama.cpp's multimodal CLI (`llama-mtmd-cli --mmproj`), which runs on
this Intel x86 CPU via AVX SIMD and sidesteps the torch==2.2.2 / transformers 4.x version trap entirely.

Manifest JSON:  [{"id": "...", "image": "/abs/path.jpg", "question": "..."}]
Output JSON:    [{"id": "...", "answer": "..."}]

Fully offline. Paths come from env vars (so the Java side needs no backend-specific args):
    LLAMACPP_MTMD_BIN    path to the llama.cpp `llama-mtmd-cli` binary (required)
    LLAMACPP_VLM_MODEL   path to the VLM GGUF weights, e.g. Qwen2.5-VL-3B-Instruct-Q4_K_M.gguf (required)
    LLAMACPP_VLM_MMPROJ  path to the matching multimodal projector GGUF (mmproj-*.gguf) (required)
    LLAMACPP_VLM_NTOKENS max tokens to generate per answer (optional, default 64)

Per-frame invocation keeps the worker simple and serverless (no loopback). The model is memory-mapped, so
after the first call the OS page cache keeps repeated loads cheap. For heavy batches a persistent
`llama-server` backend would be faster — a documented future optimisation, not required here.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys


def _clean(text: str) -> str:
    # The CLI echoes the chat-templated conversation; the real answer is the LAST assistant turn.
    text = text.replace("\x1b[0m", "")
    marker = "assistant"
    idx = text.rfind(marker)
    if idx >= 0:
        text = text[idx + len(marker):]
    text = re.sub(r"<\|[^>]*\|>", " ", text)          # strip chat control tokens like <|im_end|>
    text = text.replace("[end of text]", " ")
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    # Drop any leaked log/timing lines; keep the natural-language answer.
    lines = [ln for ln in lines
             if not re.match(r"^(llama_|main:|mtmd_|clip_|ggml_|build:|load|encoding|decoding|\d[\d.:]*\s)",
                             ln, re.IGNORECASE)]
    return " ".join(lines).strip()


def _caption(binary: str, model: str, mmproj: str, image: str, question: str, ntokens: int) -> str:
    command = [
        binary, "-m", model, "--mmproj", mmproj, "--image", image,
        "-p", question, "--temp", "0", "-n", str(ntokens), "-t", "4",
        # Force CPU: this machine's integrated GPU times out (Metal command-buffer) on the vision encoder.
        "-ngl", "0", "--no-mmproj-offload",
    ]
    proc = subprocess.run(command, capture_output=True, text=True, timeout=900)
    if proc.returncode != 0:
        sys.stderr.write("vision_caption_llamacpp: cli failed (%s): %s\n" % (proc.returncode, proc.stderr[-500:]))
        return ""
    return _clean(proc.stdout)


def main() -> int:
    parser = argparse.ArgumentParser(description="Local VLM frame captioner (llama.cpp / GGUF)")
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    binary = os.environ.get("LLAMACPP_MTMD_BIN", "")
    model = os.environ.get("LLAMACPP_VLM_MODEL", "")
    mmproj = os.environ.get("LLAMACPP_VLM_MMPROJ", "")
    ntokens = int(os.environ.get("LLAMACPP_VLM_NTOKENS", "64"))
    missing = [name for name, val in
               (("LLAMACPP_MTMD_BIN", binary), ("LLAMACPP_VLM_MODEL", model), ("LLAMACPP_VLM_MMPROJ", mmproj))
               if not val or not os.path.exists(val)]
    if missing:
        sys.stderr.write("vision_caption_llamacpp: missing/invalid env paths: %s\n" % ", ".join(missing))
        return 2

    with open(args.manifest, "r", encoding="utf-8") as handle:
        items = json.load(handle)

    results = []
    for item in items:
        answer = ""
        try:
            answer = _caption(binary, model, mmproj, item["image"], item["question"], ntokens)
        except Exception as exc:  # a single bad frame must never fail the batch
            sys.stderr.write("vision_caption_llamacpp: frame %s failed: %s\n" % (item.get("id"), exc))
        results.append({"id": item.get("id"), "answer": answer})

    os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
    with open(args.output, "w", encoding="utf-8") as handle:
        json.dump(results, handle, ensure_ascii=False, indent=2)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())