#!/usr/bin/env python3
"""Render Miril-DroneVLM-2B-2 responses beside one image."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from PIL import Image

from inference import load_model
from video_overlay import find_visual_payload, infer_responses, initial_center, render


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--prompt", action="append", required=True)
    parser.add_argument("--model-id", default="MirilAI/Miril-DroneVLM-2B-2")
    parser.add_argument("--processor-id")
    parser.add_argument("--max-new-tokens", type=int, default=384)
    parser.add_argument("--load-4bit", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    prompts = [prompt.strip() for prompt in args.prompt if prompt.strip()][:3]
    if not prompts:
        raise ValueError("at least one non-empty --prompt is required")
    model, processor = load_model(
        args.model_id,
        args.processor_id or args.model_id,
        args.load_4bit,
    )
    image = Image.open(args.image).convert("RGB")
    responses = infer_responses(
        model,
        processor,
        image,
        prompts,
        args.max_new_tokens,
    )
    precise, coarse = find_visual_payload(responses)
    center = initial_center(precise, image.width, image.height)
    rendered = render(
        image,
        responses,
        tracked_center=center,
        precise_payload=precise,
        coarse_payload=coarse,
        reveal_index=12,
    )
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    rendered.save(output)
    output.with_suffix(output.suffix + ".json").write_text(
        json.dumps(
            {
                "model_id": args.model_id,
                "image": args.image,
                "output": str(output),
                "responses": [
                    {
                        "prompt": response.prompt,
                        "payload": response.payload,
                        "validation_errors": response.errors,
                    }
                    for response in responses
                ],
            },
            indent=2,
            ensure_ascii=False,
        )
        + "\n",
        encoding="utf-8",
    )
    return 0


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