#!/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, 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. Two backends, selected by env:

  SERVER mode (fast — model loaded ONCE):     set LLAMACPP_SERVER_BIN to the `llama-server` binary.
  CLI mode  (simple — reloads per frame):     set LLAMACPP_MTMD_BIN to the `llama-mtmd-cli` binary.
                                              (server mode wins if both are set.)

  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)
  LLAMACPP_SERVER_PORT  loopback port for server mode (optional, default 8123)

CPU is forced (`-ngl 0 --no-mmproj-offload`): this machine's integrated GPU times out on the vision encoder
(Metal command-buffer timeout). Server mode keeps the model resident and POSTs base64 frames to the
OpenAI-compatible /v1/chat/completions endpoint on 127.0.0.1 — far faster than reloading per frame.
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import re
import subprocess
import sys
import time
import urllib.error
import urllib.request


def _clean(text: str) -> str:
    if not text:
        return ""
    text = text.replace("\x1b[0m", "")
    # CLI mode echoes the chat-templated conversation; take the LAST assistant turn. (Server mode returns
    # clean content already, in which case there is no 'assistant' marker and the whole string is kept.)
    idx = text.rfind("assistant")
    if idx >= 0:
        text = text[idx + len("assistant"):]
    text = re.sub(r"<\|[^>]*\|>", " ", text).replace("[end of text]", " ")
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    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()


# --- CLI mode (per-frame reload) --------------------------------------------------------------------------

def _caption_cli(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",
        "-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)


# --- SERVER mode (model loaded once) ----------------------------------------------------------------------

def _start_server(binary: str, model: str, mmproj: str, port: int):
    command = [
        binary, "-m", model, "--mmproj", mmproj, "--host", "127.0.0.1", "--port", str(port),
        "-ngl", "0", "--no-mmproj-offload", "-t", "4", "-c", "4096",
    ]
    proc = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    # Wait for readiness (model load on CPU can take a while).
    deadline = time.time() + 300
    while time.time() < deadline:
        if proc.poll() is not None:
            raise RuntimeError("llama-server exited during startup (code %s)" % proc.returncode)
        try:
            with urllib.request.urlopen("http://127.0.0.1:%d/health" % port, timeout=3) as resp:
                if resp.status == 200:
                    return proc
        except Exception:
            time.sleep(2)
    proc.terminate()
    raise RuntimeError("llama-server did not become healthy within timeout")


def _caption_server(port: int, image: str, question: str, ntokens: int) -> str:
    with open(image, "rb") as fh:
        b64 = base64.b64encode(fh.read()).decode("ascii")
    payload = {
        "messages": [{"role": "user", "content": [
            {"type": "text", "text": question},
            {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64}},
        ]}],
        "temperature": 0,
        "max_tokens": ntokens,
    }
    req = urllib.request.Request(
        "http://127.0.0.1:%d/v1/chat/completions" % port,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=900) as resp:
        body = json.loads(resp.read().decode("utf-8"))
    return _clean(body["choices"][0]["message"]["content"])


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()

    model = os.environ.get("LLAMACPP_VLM_MODEL", "")
    mmproj = os.environ.get("LLAMACPP_VLM_MMPROJ", "")
    ntokens = int(os.environ.get("LLAMACPP_VLM_NTOKENS", "64"))
    server_bin = os.environ.get("LLAMACPP_SERVER_BIN", "")
    cli_bin = os.environ.get("LLAMACPP_MTMD_BIN", "")
    port = int(os.environ.get("LLAMACPP_SERVER_PORT", "8123"))

    for name, val in (("LLAMACPP_VLM_MODEL", model), ("LLAMACPP_VLM_MMPROJ", mmproj)):
        if not val or not os.path.exists(val):
            sys.stderr.write("vision_caption_llamacpp: missing/invalid %s\n" % name)
            return 2
    use_server = bool(server_bin and os.path.exists(server_bin))
    if not use_server and not (cli_bin and os.path.exists(cli_bin)):
        sys.stderr.write("vision_caption_llamacpp: set LLAMACPP_SERVER_BIN or LLAMACPP_MTMD_BIN to a real path\n")
        return 2

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

    server = None
    results = []
    try:
        if use_server:
            server = _start_server(server_bin, model, mmproj, port)
        for item in items:
            answer = ""
            try:
                if use_server:
                    answer = _caption_server(port, item["image"], item["question"], ntokens)
                else:
                    answer = _caption_cli(cli_bin, 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})
    finally:
        if server is not None:
            server.terminate()
            try:
                server.wait(timeout=10)
            except Exception:
                server.kill()

    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())