from pathlib import Path
from io import BytesIO
import argparse
import requests
import torch
from PIL import Image

from transformers import AutoProcessor

try:
    from optimum.intel import OVModelForVisualCausalLM
except ImportError:
    from optimum.intel.openvino import OVModelForVisualCausalLM


def load_image(path_or_url: str) -> Image.Image:
    path_or_url = str(path_or_url)

    if path_or_url.startswith(("http://", "https://")):
        response = requests.get(path_or_url, timeout=30)
        response.raise_for_status()
        return Image.open(BytesIO(response.content)).convert("RGB")

    return Image.open(path_or_url).convert("RGB")


def main() -> None:
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--model-dir",
        default=".",
        help="Path to the OpenVINO-converted Qwen3.6 model directory.",
    )

    parser.add_argument(
        "--device",
        default="CPU",
        help="OpenVINO device, for example CPU, GPU, GPU.0, or AUTO.",
    )

    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=64,
        help="Maximum number of newly generated tokens.",
    )

    parser.add_argument(
        "--image",
        default="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG",
        help="Image path or image URL.",
    )

    parser.add_argument(
        "--prompt",
        default="What animal is on the candy? Answer in one sentence.",
        help="Prompt for the image.",
    )

    args = parser.parse_args()

    model_dir = Path(args.model_dir)

    if not model_dir.exists():
        raise FileNotFoundError(f"Model directory not found: {model_dir}")

    print("Using model:", model_dir)
    print("Using device:", args.device)

    processor = AutoProcessor.from_pretrained(
        model_dir,
        trust_remote_code=True,
    )

    model = OVModelForVisualCausalLM.from_pretrained(
        model_dir,
        device=args.device,
        trust_remote_code=True,
    )

    image = load_image(args.image)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": args.prompt},
            ],
        }
    ]

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    )

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=args.max_new_tokens,
            do_sample=False,
        )

    prompt_len = inputs["input_ids"].shape[-1]
    generated_ids = output_ids[0][prompt_len:]

    answer = processor.decode(
        generated_ids,
        skip_special_tokens=True,
    ).strip()

    print("\nGenerated answer:")
    print(answer)


if __name__ == "__main__":
    main()