#!/usr/bin/env python3
"""Minimal batch inference CLI for the LocateAnything-3B release code.

Examples:
  python batch_infer.py --model /path/to/LocateAnything-3B --attn sdpa \
    --image demo.jpg --query "person</c>car"

  python batch_infer.py --requests requests.jsonl --batch-size 16 --attn la_flash

Each JSONL request should contain {"image": "/path/to.jpg", "query": "person</c>car"}.
"""
import argparse
import json
import os
from pathlib import Path

from PIL import Image


def _attn_arg(value):
    mode = (value or "sdpa").strip().lower().replace("-", "_")
    aliases = {
        "": "sdpa",
        "manual": "eager",
        "torch": "eager",
        "torch_eager": "eager",
        "torch_sdpa": "sdpa",
        "flash": "la_flash",
        "la_flash": "la_flash",
        "kernel": "la_flash",
        "cuda": "la_flash",
        "range": "la_flash",
        "range_attention": "la_flash",
    }
    mode = aliases.get(mode, mode)
    if mode not in {"sdpa", "eager", "magi", "la_flash"}:
        raise argparse.ArgumentTypeError(
            f"--attn must be one of sdpa, eager, magi, la_flash; got {value!r}"
        )
    return mode


def _load_requests(args):
    requests = []
    if args.requests:
        with open(args.requests, "r", encoding="utf-8") as f:
            for line in f:
                if not line.strip():
                    continue
                row = json.loads(line)
                requests.append((row["image"], row["query"]))
    if args.image or args.query:
        if len(args.image or []) != len(args.query or []):
            raise ValueError("--image and --query must appear the same number of times")
        requests.extend(zip(args.image, args.query))
    if not requests:
        raise ValueError("provide --requests JSONL or at least one --image/--query pair")
    return requests


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--requests", help="JSONL file with image/query fields")
    ap.add_argument("--image", action="append", help="Image path; repeat with --query")
    ap.add_argument("--query", action="append", help="Category query, e.g. person</c>car")
    ap.add_argument("--model", default=os.environ.get("LA_FLASH_MODEL", "nvidia/LocateAnything-3B"))
    ap.add_argument("--attn", type=_attn_arg, default=os.environ.get("LA_FLASH_ATTN", "sdpa"),
                    help="LLM attention backend: sdpa, eager, magi, or la_flash")
    ap.add_argument("--vision-attn", default=os.environ.get("LA_FLASH_VISION_ATTN", "auto"),
                    choices=["auto", "flash_attention_2", "sdpa", "eager"])
    ap.add_argument("--batch-size", type=int, default=1)
    ap.add_argument("--scheduler", default=os.environ.get("LA_FLASH_HYBRID_SCHEDULER", "eager"),
                    choices=["eager", "hold_ar", "ar_first", "pipeline", "adaptive"])
    ap.add_argument("--group-size", type=int, default=int(os.environ.get("LA_FLASH_HYBRID_GROUP_SIZE", "0")))
    ap.add_argument("--max-new-tokens", type=int, default=2048)
    ap.add_argument("--temperature", type=float, default=0.7)
    ap.add_argument("--top-p", type=float, default=0.9)
    ap.add_argument("--top-k", type=int, default=0)
    ap.add_argument("--repetition-penalty", type=float, default=1.1)
    ap.add_argument("--strict-attn", action="store_true",
                    help="Fail instead of falling back to SDPA if magi/la_flash is unavailable")
    ap.add_argument("--out", default="", help="Optional output JSONL path; stdout if omitted")
    args = ap.parse_args()
    args.attn = _attn_arg(args.attn)

    os.environ["LA_FLASH_MODEL"] = args.model
    os.environ["LA_FLASH_ATTN"] = args.attn
    os.environ["LA_FLASH_VISION_ATTN"] = args.vision_attn
    os.environ["LA_FLASH_HYBRID_SCHEDULER"] = args.scheduler
    os.environ["LA_FLASH_HYBRID_GROUP_SIZE"] = str(args.group_size)
    if args.strict_attn:
        os.environ["LA_FLASH_STRICT_ATTN"] = "1"

    from batch_utils import generate_batch_hybrid, get_last_hybrid_stats, load
    from batch_utils.hybrid_runtime import load_pil

    requests = _load_requests(args)
    load()

    writer = open(args.out, "w", encoding="utf-8") if args.out else None
    try:
        for start in range(0, len(requests), max(1, args.batch_size)):
            chunk = requests[start:start + max(1, args.batch_size)]
            pairs = [(load_pil(image), query) for image, query in chunk]
            texts = generate_batch_hybrid(
                pairs,
                temperature=args.temperature,
                top_p=None if args.top_p < 0 else args.top_p,
                top_k=None if args.top_k <= 0 else args.top_k,
                repetition_penalty=args.repetition_penalty,
                max_new_tokens=args.max_new_tokens,
                scheduler=args.scheduler,
                group_size=args.group_size,
            )
            stats = get_last_hybrid_stats()
            for (image, query), text in zip(chunk, texts):
                row = {"image": str(Path(image)), "query": query, "raw_response": text, "stats": stats}
                line = json.dumps(row, ensure_ascii=False)
                if writer:
                    writer.write(line + "\n")
                else:
                    print(line, flush=True)
    finally:
        if writer:
            writer.close()


if __name__ == "__main__":
    main()