#!/usr/bin/env python3
"""Compile Qwen2.5-VL's vision tower or text decoder to mblt through ``mblt_compile``.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import requests
import torch
from PIL import Image
from transformers import AutoProcessor

from qbcompiler.frontend import mblt_compile
from qbcompiler.model_dict.parser.backend.torch.input_capture import (
    capture_forward_inputs,
)
from qbcompiler.model_dict.parser.patcher.parts import load_for_part, prepare_part

DEFAULT_MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"
DEFAULT_IMAGE_URL = "http://images.cocodataset.org/val2017/000000039769.jpg"
DEFAULT_PROMPT = "Please describe the image explicitly."
DEFAULT_TARGET_DEVICE = "aries-rb"
DEFAULT_IMAGE_SIZE = (224, 224)

# Named against the feed: cache_position is 4.x-only, rope_deltas may be absent.
LANGUAGE_SEQ_INPUTS = ("inputs_embeds", "cache_position", "rope_deltas")
LANGUAGE_SEQ_AXIS = {"inputs_embeds": [-2], "cache_position": [-1], "rope_deltas": [-1]}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Compile Qwen2.5-VL vision/language.")
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID, help="HF model id.")
    parser.add_argument("--base-path", default=".", help="Output directory.")
    parser.add_argument(
        "--part",
        choices=["vision", "language"],
        default="language",
        help="Which part to compile.",
    )
    parser.add_argument("--image-url", default=DEFAULT_IMAGE_URL, help="Sample image.")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="Sample prompt.")
    parser.add_argument("--image-size", type=int, nargs=2, default=DEFAULT_IMAGE_SIZE)
    parser.add_argument("--target-device", default=DEFAULT_TARGET_DEVICE)
    parser.add_argument("--max-new-tokens", type=int, default=1)
    parser.add_argument("--torch-device", default="cpu", help="auto|cuda|cpu")
    parser.add_argument(
        "--dtype", default="bfloat16", choices=["float32", "float16", "bfloat16"]
    )
    parser.add_argument(
        "--static", action="store_true", default=False, help="No dynamic axes."
    )
    return parser.parse_args()


def resolve_device(name: str) -> torch.device:
    if name == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if name == "cuda" and not torch.cuda.is_available():
        return torch.device("cpu")
    return torch.device(name)


def resolve_dtype(name: str) -> torch.dtype:
    if name == "float16":
        return torch.float16
    if name == "bfloat16":
        return torch.bfloat16
    return torch.float32


def build_inputs(processor, image, prompt, device, dtype):
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]
    return processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(device=device, dtype=dtype)


def main() -> None:
    args = parse_args()
    base_path = Path(args.base_path)
    base_path.mkdir(parents=True, exist_ok=True)
    save_name = args.model_id.replace("/", "_")
    mblt_path = base_path / f"{save_name}_{args.part}.mblt"

    device = resolve_device(args.torch_device)
    dtype = resolve_dtype(args.dtype)

    processor = AutoProcessor.from_pretrained(args.model_id)
    image = Image.open(requests.get(args.image_url, stream=True).raw).convert("RGB")
    image = image.resize(tuple(args.image_size))

    model = load_for_part(args.model_id, args.part, dtype=dtype, device=device)
    inputs = build_inputs(processor, image, args.prompt, device, dtype)

    capture_target = prepare_part(model, args.part)
    with capture_forward_inputs(capture_target, to_cpu=False) as feed_dict:
        model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False)
    feed_dict = dict(feed_dict)

    if args.part == "vision" or args.static:
        dynamic_axes = None
    else:
        dynamic_axes = {
            name: LANGUAGE_SEQ_AXIS[name]
            for name in LANGUAGE_SEQ_INPUTS
            if name in feed_dict
        }

    mblt_compile(
        model=model,
        model_part=args.part,
        backend="torch",
        target_device=args.target_device,
        mblt_save_path=str(mblt_path),
        feed_dict=feed_dict,
        dynamic_axes=dynamic_axes,
    )
    print(f"write: {mblt_path}")


if __name__ == "__main__":
    main()