#!/usr/bin/env python3
"""Tier-2 vision captioner for the highlight director.

Loads a local vision-language model (moondream2) ONCE and answers a question for each frame listed in a
manifest, writing the answers to a JSON file. Runs fully offline from the local Hugging Face cache — it
performs no network access (HF_HUB_OFFLINE is forced on).

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

import argparse
import json
import os
import sys


def main() -> int:
    parser = argparse.ArgumentParser(description="Local VLM frame captioner (moondream2)")
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--model-id", default="vikhyatk/moondream2")
    parser.add_argument("--revision", default="2024-08-26")
    args = parser.parse_args()

    # Hard offline: never reach the network. The model must already be in the local cache.
    os.environ["HF_HUB_OFFLINE"] = "1"
    os.environ["TRANSFORMERS_OFFLINE"] = "1"
    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

    try:
        from PIL import Image
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except Exception as exc:  # pragma: no cover - dependency guard
        print(f"vision_caption: missing dependency: {exc}", file=sys.stderr)
        return 3

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

    try:
        model = AutoModelForCausalLM.from_pretrained(
            args.model_id, revision=args.revision, trust_remote_code=True, local_files_only=True)
        tokenizer = AutoTokenizer.from_pretrained(
            args.model_id, revision=args.revision, local_files_only=True)
        model.eval()
    except Exception as exc:
        print(f"vision_caption: unable to load local model: {exc}", file=sys.stderr)
        return 2

    results = []
    for item in items:
        try:
            image = Image.open(item["image"]).convert("RGB")
            encoded = model.encode_image(image)
            answer = model.answer_question(encoded, item["question"], tokenizer).strip()
        except Exception as exc:  # keep going; a single bad frame must not fail the batch
            print(f"vision_caption: frame {item.get('id')} failed: {exc}", file=sys.stderr)
            answer = ""
        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())