#!/usr/bin/env python3
"""Apple Silicon inference for Miril-DroneVLM-2B-2 MLX variants."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from inference import dispatch_summary, parse_bare_json
from router_contract import TYPED_JSON_ROUTER_SYSTEM_PROMPT


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


def main() -> int:
    from mlx_vlm import generate, load

    args = parse_args()
    model, processor = load(args.model_id, strict=True)
    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": TYPED_JSON_ROUTER_SYSTEM_PROMPT}],
        },
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": args.prompt},
            ],
        },
    ]
    rendered_prompt = processor.apply_chat_template(
        messages, add_generation_prompt=True
    )
    generated = generate(
        model,
        processor,
        prompt=rendered_prompt,
        image=args.image,
        max_tokens=args.max_tokens,
        temperature=0.0,
        verbose=False,
    )
    raw = getattr(generated, "text", str(generated)).strip()
    payload, errors = parse_bare_json(raw)
    result = {
        "ok": payload is not None,
        "model_id": args.model_id,
        "prompt": args.prompt,
        "raw": raw,
        "payload": payload,
        "validation_errors": errors,
        "dispatch": dispatch_summary(payload) if payload is not None else None,
    }
    rendered = json.dumps(result, indent=2, ensure_ascii=False)
    print(rendered)
    if args.output_json:
        output_path = Path(args.output_json)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(rendered + "\n", encoding="utf-8")
    return 0 if not errors else 2


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