from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from transformers import AutoModelForMultimodalLM, AutoProcessor

MODEL_TITLE = "LumynaX Infused Gemma E4B Model"
SUPPORTED_MODALITIES = ('text', 'image', 'audio')
DEFAULT_ENABLE_THINKING = True


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            f"Run a local Gemma E4B quickstart for {MODEL_TITLE}. "
            f"Supported modalities: text, image, audio."
        )
    )
    parser.add_argument("--mode", choices=["text", "image", "audio"], default="text")
    parser.add_argument(
        "--prompt",
        default="Explain in two short bullet points what this local package is.",
        help="Text instruction to send to the model.",
    )
    parser.add_argument("--image", default="", help="Local image path or image URL for --mode image.")
    parser.add_argument("--audio", default="", help="Local audio path or audio URL for --mode audio.")
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument(
        "--thinking",
        action=argparse.BooleanOptionalAction,
        default=DEFAULT_ENABLE_THINKING,
        help="Enable Gemma reasoning mode.",
    )
    return parser


def _message_content(args: argparse.Namespace) -> list[dict[str, str]]:
    if args.mode == "text":
        return [{"type": "text", "text": args.prompt}]
    if args.mode == "image":
        if not args.image:
            raise SystemExit("--image is required when --mode image is used.")
        image_ref = args.image.strip()
        return [
            {"type": "image", "url": image_ref},
            {"type": "text", "text": args.prompt},
        ]
    if not args.audio:
        raise SystemExit("--audio is required when --mode audio is used.")
    audio_ref = args.audio.strip()
    return [
        {"type": "audio", "audio": audio_ref},
        {"type": "text", "text": args.prompt},
    ]


def main() -> None:
    args = _build_parser().parse_args()
    model_dir = Path(__file__).resolve().parent / "merged_model"
    if not model_dir.exists():
        raise SystemExit(f"Expected merged_model/ at {model_dir}")

    processor = AutoProcessor.from_pretrained(model_dir)
    model = AutoModelForMultimodalLM.from_pretrained(
        model_dir,
        dtype="auto",
        device_map="auto",
    )
    messages = [
        {
            "role": "user",
            "content": _message_content(args),
        },
    ]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        add_generation_prompt=True,
        enable_thinking=args.thinking,
    ).to(model.device)
    input_len = inputs["input_ids"].shape[-1]

    with torch.inference_mode():
        outputs = model.generate(**inputs, max_new_tokens=args.max_new_tokens)

    response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)
    parsed = processor.parse_response(response) if hasattr(processor, "parse_response") else response
    if isinstance(parsed, str):
        print(parsed)
        return
    print(json.dumps(parsed, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()