#!/usr/bin/env python3
import argparse
import os
import torch

from qwen_vl_utils import process_vision_info
from transformers import (
    AutoProcessor,
    AutoTokenizer,
    AutoModelForImageTextToText,
)


DEFAULT_SYSTEM_PROMPT = (
    "你是一个图像鉴伪专家，擅长结合视觉，文字结合伪造特征分析手段鉴别输入图像的真假。"
    "分析过程中，你会逐步分析，抽丝剥茧，找到图像伪造的蛛丝马迹，"
    "最终给出专业的鉴别结果及分析。"
)

DEFAULT_USER_PROMPT = (
    "请分析这张文档图片是否存在伪造或篡改风险，并输出一份专业、精炼、准确的防伪分析报告。"
)


def build_messages(image_path: str, system_prompt: str, user_prompt: str):
    return [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": system_prompt,
                }
            ],
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image_path,
                },
                {
                    "type": "text",
                    "text": user_prompt,
                },
            ],
        },
    ]


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--model-name",
        type=str,
        default="vankey/DocShield-9B",
    )
    parser.add_argument(
        "--image",
        type=str,
        required=True,
        help="Input image path.",
    )
    parser.add_argument(
        "--prompt",
        type=str,
        default=DEFAULT_USER_PROMPT,
    )
    parser.add_argument(
        "--system-prompt",
        type=str,
        default=DEFAULT_SYSTEM_PROMPT,
    )

    thinking_group = parser.add_mutually_exclusive_group()
    thinking_group.add_argument(
        "--thinking",
        action="store_true",
        help="Enable Qwen3.5 thinking mode.",
    )
    thinking_group.add_argument(
        "--no-thinking",
        action="store_true",
        help="Disable Qwen3.5 thinking mode.",
    )

    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=1024,
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
    )
    parser.add_argument(
        "--top-p",
        type=float,
        default=0.8,
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=20,
    )
    parser.add_argument(
        "--do-sample",
        action="store_true",
        help="Use sampling. If not set, greedy decoding is used.",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
    )
    parser.add_argument(
        "--dtype",
        type=str,
        default="bf16",
        choices=["bf16", "fp16", "fp32"],
    )

    return parser.parse_args()


def get_torch_dtype(dtype: str):
    if dtype == "bf16":
        return torch.bfloat16
    if dtype == "fp16":
        return torch.float16
    return torch.float32


def main():
    args = parse_args()

    if not os.path.exists(args.image):
        raise FileNotFoundError(f"Image not found: {args.image}")

    enable_thinking = False
    if args.thinking:
        enable_thinking = True
    if args.no_thinking:
        enable_thinking = False

    torch_dtype = get_torch_dtype(args.dtype)

    print("=" * 100)
    print("Model:", args.model_name)
    print("Image:", args.image)
    print("Prompt:", args.prompt)
    print("Enable thinking:", enable_thinking)
    print("Max new tokens:", args.max_new_tokens)
    print("dtype:", args.dtype)
    print("=" * 100)

    tokenizer = AutoTokenizer.from_pretrained(
        args.model_name,
        use_fast=True,
        trust_remote_code=True,
    )

    processor = AutoProcessor.from_pretrained(
        args.model_name,
        trust_remote_code=True,
    )

    model = AutoModelForImageTextToText.from_pretrained(
        args.model_name,
        torch_dtype=torch_dtype,
        trust_remote_code=True,
        device_map="auto",
    )

    model.eval()

    messages = build_messages(
        image_path=args.image,
        system_prompt=args.system_prompt,
        user_prompt=args.prompt,
    )

    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=enable_thinking,
    )

    print("\n" + "=" * 100)
    print("Rendered prompt preview:")
    print(text[:2000])
    print("=" * 100 + "\n")

    image_inputs, video_inputs = process_vision_info(messages)

    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )

    inputs = inputs.to(model.device)

    print("input_ids shape:", inputs["input_ids"].shape)
    if "pixel_values" in inputs:
        print("pixel_values shape:", inputs["pixel_values"].shape)
    if "image_grid_thw" in inputs:
        print("image_grid_thw:", inputs["image_grid_thw"])

    generation_kwargs = {
        "max_new_tokens": args.max_new_tokens,
    }

    if args.do_sample:
        generation_kwargs.update(
            {
                "do_sample": True,
                "temperature": args.temperature,
                "top_p": args.top_p,
                "top_k": args.top_k,
            }
        )
    else:
        generation_kwargs.update(
            {
                "do_sample": False,
            }
        )

    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            **generation_kwargs,
        )

    generated_ids_trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
    ]

    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=False,
        clean_up_tokenization_spaces=False,
    )[0]

    print("\n" + "=" * 100)
    print("Model output:")
    print(output_text)
    print("=" * 100)


if __name__ == "__main__":
    main()