#!/usr/bin/env python3
"""Image + text generation with Piko-9b.

    python examples/inference_multimodal.py --image receipt.png \
        --prompt "Give the merchant and total as JSON."

Read reports/inference_validation.json before relying on this path. The vision
tower in this checkpoint was copied verbatim from Qwen/Qwen3.5-9B and was never
trained or re-aligned against Piko's fine-tuned language backbone.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from urllib.parse import urlparse

import torch
from _common import add_common_arguments, generation_kwargs, load_model, strip_reasoning

DEFAULT_SYSTEM = (
    "You are Piko-9, an AI assistant. Examine the supplied image, answer accurately, "
    "read visible text when relevant, and do not invent details the image does not show."
)


def resolve_image(reference: str) -> str:
    """Accept a local path or an http(s) URL; fail early and clearly otherwise."""
    parsed = urlparse(reference)
    if parsed.scheme in ("http", "https"):
        return reference
    path = Path(reference).expanduser()
    if not path.is_file():
        sys.exit(f"Image not found: {path}")
    if path.suffix.lower() not in {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}:
        sys.exit(f"Unsupported image type: {path.suffix}")
    try:
        from PIL import Image

        with Image.open(path) as image:
            image.verify()
    except ImportError:
        sys.exit("Pillow is required: pip install pillow")
    except Exception as exc:  # noqa: BLE001
        sys.exit(f"Could not read {path} as an image: {exc}")
    return str(path.resolve())


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    add_common_arguments(parser)
    parser.add_argument(
        "--image", required=True, action="append", help="Path or URL. Repeat for multiple images."
    )
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--system", default=DEFAULT_SYSTEM)
    args = parser.parse_args()

    images = [resolve_image(reference) for reference in args.image]
    model, processor = load_model(args.model, args.quantization, args.dtype, args.revision)

    content: list[dict[str, str]] = [{"type": "image", "url": image} for image in images]
    content.append({"type": "text", "text": args.prompt})

    messages = []
    if args.system:
        messages.append({"role": "system", "content": args.system})
    messages.append({"role": "user", "content": content})

    try:
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(model.device)
    except ImportError as exc:
        if "orchvision" in str(exc):
            sys.exit("Image input needs torchvision: pip install torchvision")
        raise

    with torch.inference_mode():
        output = model.generate(**inputs, **generation_kwargs(args))

    text = processor.decode(
        output[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
    ).strip()

    print(text if args.show_reasoning else strip_reasoning(text))


if __name__ == "__main__":
    main()