import argparse
import json
import os

import cv2
import torch
from PIL import Image
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
from qwen_vl_utils import process_vision_info

MODEL_NAME = "DocShield-7B"
MODEL_PATH = "vankey/DocShield-7B"
IMG_W, IMG_H = 1344, 896

SYSTEM_PROMPT = (
    "You are a top-tier image forgery analysis expert, specialized in forensic-level "
    "image and text correlation analysis. Based on the input, produce a concise, "
    "professional, and accurate forgery analysis report."
)
USER_PROMPT = "Please analyze whether this image is forged and provide an analysis report."

GEN_KWARGS = dict(
    do_sample=True,
    temperature=1.0,
    top_p=1.0,
    top_k=0,
    repetition_penalty=1.0,
)


def load_model():
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.float32,
        attn_implementation="eager",
        device_map="auto",
    )
    model.eval()
    return model


def load_image(image_path):
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Cannot read image: {image_path}")
    image = cv2.resize(image, (IMG_W, IMG_H))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return Image.fromarray(image)


def build_messages(image_path):
    img = load_image(image_path)
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": USER_PROMPT},
            ],
        },
    ]
    return messages


def infer(model, processor, image_path, max_new_tokens=8192):
    messages = build_messages(image_path)
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.no_grad():
        output_ids = model.generate(**inputs, **GEN_KWARGS, max_new_tokens=max_new_tokens)

    generated_ids = output_ids[:, inputs["input_ids"].shape[1]:]
    result = processor.batch_decode(
        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    return result.strip()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True)
    parser.add_argument("--output", default=None)
    parser.add_argument("--max-new-tokens", type=int, default=8192)
    args = parser.parse_args()

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

    print(f"[{MODEL_NAME}] loading model: {MODEL_PATH}")
    processor = AutoProcessor.from_pretrained(MODEL_PATH)
    model = load_model()

    print(f"[{MODEL_NAME}] inferring image: {args.image}")
    answer = infer(model, processor, args.image, max_new_tokens=args.max_new_tokens)

    output_data = {"model": MODEL_NAME, "image": args.image, "answer": answer}
    output_path = args.output or os.path.join(
        os.getcwd(), f"{os.path.splitext(os.path.basename(args.image))[0]}.json"
    )
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(output_data, f, ensure_ascii=False, indent=4)

    print("\n--- result ---")
    print(answer)
    print(f"\n--- saved to: {output_path} ---")


if __name__ == "__main__":
    main()