# -*- coding: utf-8 -*-
"""
使用方法:
python rule_14_vllm.py \
  --input_dir "/path/to/your/images" \
  --model_path "/path/to/your/Qwen2.5-VL-7B-Instruct"
"""

import os
import re
import json
import argparse
import multiprocessing
from pathlib import Path
from typing import Dict, List, Optional, Any
from tqdm import tqdm
from PIL import Image, ImageFile
from transformers import AutoProcessor
from vllm import LLM, SamplingParams

os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'


ImageFile.LOAD_TRUNCATED_IMAGES = True


IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}



SYS_PROMPT_TEXT="""You are a highly critical Senior Art Director. Your job is to flag "Low-Quality / Amateur" advertising designs. Crucial Context: You must distinguish between "Aggressive E-commerce Marketing" (Professional) and "Amateur Sloppiness" (Violation). High-resolution assets, standard platform badges, and professional 3D renders are SAFE.

INPUT: One image and one natural-language question about design aesthetic and text harmony.

YOUR TASK:

Determine if the image is a VIOLATION (Unsuitable) or SAFE (Suitable) based on the criteria below.

Output a JSON object containing a rigorous Chain-of-Thought ("think") and a precise classification label ("answer").

OUTPUT FORMAT: Return EXACTLY two blocks, no extra text: <think>Detailed reasoning evaluating font effects, background integration, and aesthetic consistency...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Text Design Harmony"}</answer>

========================================= STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
The "WordArt" Effect (廉价特效):

Technical Failure: ONLY flag if text is pixelated, jagged, or uses 1990s-style rainbow/neon gradients.

Distortion: Text is unprofessionally stretched or squeezed (breaking the font's aspect ratio).

Amateur Strokes: Thick, vibrating outlines that look like they were made in MS Paint, not professional design software.

Note: High-res 3D fonts, clean gold textures, and smooth gradients are PROFESSIONAL and SAFE.

Visual Clutter & Conflict (背景冲突与拼贴感):

Resolution Mismatch: A low-res/blurry graphic pasted onto a high-res photo.

Zero Integration: Elements that have NO shadows, NO lighting consistency, and look like accidental "floating" errors.

Legibility Loss: Text is truly unreadable due to background chaos without any masking.

Note: Standard UI elements (Pill buttons, Price tags, Promo badges like "百亿补贴") are INTENTIONAL overlays and are SAFE.

Inconsistent Aesthetic (风格割裂):

Flag ONLY if elements are accidentally mismatched (e.g., a hand-drawn sketch randomly appearing in a high-tech 3D render without stylistic intent).

Note: 3D mascots or cartoon characters placed on realistic backgrounds for marketing purposes are a VALID style and are SAFE.

========================================= CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
Commercial Execution: High-resolution assets, clean font edges, and professional lighting/shadows.

Platform Legitimacy: Presence of brand logos (Alipay, Taobao, Banks, China Gold) and standard e-commerce UI components.

Intentional Hierarchy: Even if the design is "loud" (Red/Gold), it is Suitable if the text is aligned and the layout is purposeful.

========================================= DECISION LOGIC
Unsuitable: If the design shows Technical Failure (pixelation, distortion, 90s-style WordArt) or looks like a non-designer's mistake.

Suitable: If the design follows Commercial Logic (Standard e-commerce banners, High-res renders, Professional marketing layouts). When in doubt, if the image looks like it's from a major App, it is SUITABLE.
"""



def collect_images(input_dir: Path) -> List[Dict[str, str]]:
    if not input_dir.exists():
        raise FileNotFoundError(f"Input directory not found: {input_dir}")
    
    files = [p for p in input_dir.iterdir() if p.is_file() and p.suffix.lower() in IMG_EXTS]
    files.sort()
    
    print(f"[Info] Found {len(files)} images in {input_dir}")
    return [{"path": str(p), "filename": p.name} for p in files]

def parse_llm_output(text: str) -> Dict[str, Any]:
    default_res = {
        "label": "Parse Error", 
        "think": "No reasoning found", 
        "raw": text
    }
    
    if not text:
        return default_res

    think_match = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
    think_content = think_match.group(1).strip() if think_match else ""

    answer_match = re.search(r'<answer>(.*?)</answer>', text, re.DOTALL)
    
    extracted_label = "Parse Error"
    
    if answer_match:
        json_str = answer_match.group(1).strip()
        try:
            data = json.loads(json_str)
            raw_ans = data.get("Answer", "")
            
            if "unsuitable" in raw_ans.lower():
                extracted_label = "Unsuitable"
            elif "suitable" in raw_ans.lower():
                extracted_label = "Suitable"
            else:
                extracted_label = raw_ans 
                
        except json.JSONDecodeError:
            if "Unsuitable" in json_str:
                extracted_label = "Unsuitable"
            elif "Suitable" in json_str:
                extracted_label = "Suitable"
    else:
        if "Unsuitable" in text:
            extracted_label = "Unsuitable"
        elif "Suitable" in text:
            extracted_label = "Suitable"

    return {
        "label": extracted_label,
        "think": think_content,
        "raw": text
    }

def prepare_vllm_inputs(batch_meta: List[Dict], processor) -> List[Dict]:
    vllm_inputs = []
    user_query = "Analyze this image against the design rules and return the JSON decision."
    
    for item in batch_meta:
        img_path = item["path"]
        try:
            image_obj = Image.open(img_path).convert("RGB")
            
            messages = [
                {"role": "system", "content": [{"type": "text", "text": SYS_PROMPT_TEXT}]},
                {"role": "user", "content": [
                    {"type": "image", "image": img_path},
                    {"type": "text", "text": user_query}
                ]}
            ]
            
            prompt_text = processor.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            
            vllm_inputs.append({
                "prompt": prompt_text,
                "multi_modal_data": {"image": image_obj}
            })
        except Exception as e:
            print(f"[Warning] Failed to load {img_path}: {e}")
            vllm_inputs.append(None)
            
    return vllm_inputs


def main():
    parser = argparse.ArgumentParser(description="AI Visual Comfort Auditor")
    parser.add_argument("--input_dir", type=str, required=True, help="Folder containing images to check")
    parser.add_argument("--model_path", type=str, required=True, help="Path to local Qwen-VL model")
    parser.add_argument("--batch_size", type=int, default=512, help="Inference batch size")
    parser.add_argument("--tp_size", type=int, default=2, help="Tensor Parallel size")
    args = parser.parse_args()

    input_path = Path(args.input_dir)
    meta_data = collect_images(input_path)
    
    if not meta_data:
        print("[Info] No images found. Exiting.")
        return
    print(f"\n[Init] Loading Model: {args.model_path}")
    
    llm = LLM(
        model=args.model_path,
        tokenizer=args.model_path,
        trust_remote_code=True,
        tensor_parallel_size=args.tp_size,
        gpu_memory_utilization=0.90, 
        max_model_len=8192, 
        enforce_eager=True,
        limit_mm_per_prompt={"image": 1}
    )
    
    processor = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True)
    
    sampling_params = SamplingParams(
        temperature=0.7, 
        max_tokens=1024,
        top_p=0.9
    )
    results = []
    print(f"\n[Run] Starting Inference on {len(meta_data)} images...")
    
    for i in tqdm(range(0, len(meta_data), args.batch_size), desc="Processing Batches"):
        batch_meta = meta_data[i : i + args.batch_size]
        batch_inputs = prepare_vllm_inputs(batch_meta, processor)
        
        valid_inputs = [inp for inp in batch_inputs if inp is not None]
        valid_indices = [idx for idx, inp in enumerate(batch_inputs) if inp is not None]
        
        if not valid_inputs:
            continue
            
        outputs = llm.generate(valid_inputs, sampling_params=sampling_params, use_tqdm=False)
        
        for local_idx, out in enumerate(outputs):
            original_meta = batch_meta[valid_indices[local_idx]]
            generated_text = out.outputs[0].text
            
            parsed = parse_llm_output(generated_text)
            
            results.append({
                "filename": original_meta["filename"],
                "path": original_meta["path"],
                "label": parsed["label"],   # Suitable / Unsuitable
                "think": parsed["think"],  
                "raw_output": generated_text
            })

    total = len(results)
    unsuitable_count = sum(1 for r in results if r["label"] == "Unsuitable")
    suitable_count = sum(1 for r in results if r["label"] == "Suitable")
    error_count = total - unsuitable_count - suitable_count

    unsuitable_rate = (unsuitable_count / total * 100) if total > 0 else 0
    suitable_rate = (suitable_count / total * 100) if total > 0 else 0

    print("\n" + "="*60)
    print(f"AUDIT REPORT FOR: {input_path.name}")
    print("="*60)
    print(f"{'Total Images':<25}: {total}")
    print("-" * 60)
    print(f"{'UNSUITABLE (Violation)':<25}: {unsuitable_count}  ({unsuitable_rate:.2f}%)")
    print(f"{'SUITABLE (Safe)':<25}: {suitable_count}  ({suitable_rate:.2f}%)")
    print(f"{'Parse Errors':<25}: {error_count}")
    print("="*60)

    output_file = input_path / f"audit_result_{input_path.name}.json"
    try:
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"\n[Done] Detailed JSON report saved to:\n-> {output_file}")
    except Exception as e:
        print(f"[Error] Could not save JSON: {e}")

if __name__ == "__main__":
    try:
        multiprocessing.set_start_method('spawn', force=True)
    except RuntimeError:
        pass
    
    main()