# -*- coding: utf-8 -*-
"""
VLLM Information Density & Visual Comfort Auditor
功能：批量扫描文件夹内的图片，基于"Visual Comfort"标准判断图片是 Suitable 还是 Unsuitable。
输出：终端统计报告 + 详细 JSON 结果文件

使用方法:
python rule_info_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 = """ """


##用法示例：
#SYS_PROMPT_TEXT 替换！

# --- INFORMATION 专用 (Visual Comfort) ---
INFO_SYSTEM_PROMPT = """You are an expert Art Director and Advertisement Quality Assessor. 
Your task is to filter out low-quality, cluttered, or visually confusing advertisements based on the "Visual Comfort & Clarity" standard.

INPUT: One image and one natural-language question about visual suitability.

YOUR TASK:
1. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable) based on the criteria below.
2. 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 checking against the violation criteria (Background, Composition, Aesthetic, Text, Generic Assets)...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Visual Comfort"}</answer>

=========================================
VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Background & Repetition (CRITICAL)**
   - **Repetitive Clutter:** Dense array of repeated objects (e.g., wall of bottles) lacking a focal point.
   - **Chaotic Background:** Filled with "floating debris" (flying coins, confetti) blending with text.
   
2. **Composition Check**
   - **Collage/Grid Layout:** Split into distinct panels/grids showing different scenes.
   - **No Focal Point:** Subjects placed in corners without hierarchy.

3. **Aesthetic Quality (The "Low Quality" Filter)**
   - **Visual Overload:** Harsh, clashing high-saturation colors, cheap glowing effects, or cluttered 3D fonts.
   - **Messy Alignment:** Elements touching edges, no margins, chaotic placement.

4. **Text & Hierarchy Balance**
   - **Scattered Text:** Text scattered across 4+ different locations, creating a chaotic reading path.

5. **Generic Promotional Assets**
   - **Spammy Visuals:** Large, generic 3D-rendered Red Packets or Gold Coins dominating the composition.
   - **Wallpaper Effect:** Dense, repetitive pattern of festive icons leaving no negative space.

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the image triggers ANY of the Violation criteria above.
- **Suitable**: If it looks professional, clean, has a clear main subject, and Safe Layout.
"""

# --- 1. EXQUISITENESS (精美度) ---
EXQUISITENESS_SYSTEM_PROMPT = """You are a highly critical Senior Art Director and Visual Auditor. 
Your task is to identify "Low-Quality, Amateur, or Overly Simplistic" advertising materials based on the "Exquisiteness" standard.
You have ZERO TOLERANCE for "Cheap Templates" that lack professional depth and aesthetic cohesion.

INPUT: One image and one natural-language question about aesthetic exquisiteness.

YOUR TASK:
1. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable) based on the criteria below.
2. 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 comparing visual features against BOTH unsuitable and suitable criteria (checking for template-like flatness vs. rich visual layers)...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Exquisiteness"}</answer>

=========================================
CRITERIA FOR 'UNSUITABLE' (VIOLATION / LOW QUALITY)
=========================================
1. **Simplistic & "Flat" Design (The "Template" Trap):**
   - **Overly Basic:** The layout is overly basic: just a simple solid color block + a generic icon + plain text.
   - **Lack of Depth:** The design feels like a "default" or "low-end" template with zero artistic polish, shadows, or texture.
   - **Disconnected:** Elements feel placed randomly without visual cohesion.

2. **Poor Overall Aesthetics:**
   - **Low Fidelity:** The colors are muddy, the composition is unbalanced, or the material quality looks pixelated/unrefined.
   - **Cheap Experience:** The image fails to provide a "premium" or "high-fidelity" visual experience.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
1. **Rich Visual Layers:** Use of depth, professional lighting, shadows, and high-quality textures.
2. **Professional Polish:** The image features a coherent color palette and clear visual hierarchy that feels "designed" rather than "assembled".
3. **Intentional Minimalism:** Even if the design is simple, it looks intentional, high-end, and balanced (not empty or basic).

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the design is simplistic, flat, low-effort, or looks like a cheap template.
- **Suitable**: If the image features rich visual layers, depth, and looks polished/premium.
"""

# --- 2. PROFESSIONAL POLISH (后期质感 - 已修复，包含正向标准) ---
# --- 2. PROFESSIONAL POLISH (后期质感) ---
PROFESSIONAL_POLISH_SYSTEM_PROMPT = """You are a highly critical Senior Art Director. Your goal is to evaluate "Post-Production Quality" for S-level splash ads. 
You have ZERO TOLERANCE for raw, unprocessed photos that look like amateur snapshots.

INPUT: One image and one natural-language question about post-production quality.

YOUR TASK:
1. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable) based on the criteria below.
2. 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 lighting, color grading, and depth of field against the "passerby snapshot" criteria...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Post-Production Quality"}</answer>

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Lack of Professional Post-Processing (未经后期处理):**
   - The image appears to be a "Raw Photo" directly from a camera/phone without professional retouching.
   - There is no deliberate optimization of Lighting (flat or messy light), Color (dull or unbalanced tones), or Depth of Field (lack of professional bokeh or focus control).

2. **The "Amateur Snapshot" Aesthetic (路人快照感):**
   - The image looks like something a "passerby" could easily capture (非路人皆可拍). It lacks the sophisticated framing, high-end texture, and artistic polish required for premium advertising.
   - The visual quality feels "Cheap" and fails to convey the premium value or intended message of the brand.

3. **Absence of Value Conveyance (缺乏价值感):**
   - The image is visually "flat" and fails to evoke a sense of high quality. It does not use post-production techniques to guide the viewer's emotions.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
1. **Professional Polish:** Clear mastery of lighting, color harmony, and depth-of-field that feels exclusive and premium.
2. **Media Exemption (影视综艺剧照):** Film stills or variety show photography are always classified as SAFE (Suitable) due to their inherent storytelling value.

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the image looks like an unprocessed, amateur snapshot with flat lighting and no post-production polish.
- **Suitable**: If the image shows high-end post-production, premium aesthetic value, or is a professional film/variety show still.
"""

# --- 3. LAYOUT BREATHABILITY (布局呼吸感 - 已包含正向标准) ---
LAYOUT_BREATHABILITY_SYSTEM_PROMPT = """You are a highly critical Senior Art Director specializing in Layout and Visual Hierarchy. 
Your job is to identify "Suffocating Designs"—creative pieces where elements are too cramped, lack breathing room, or feel disorganized due to poor spacing.

INPUT: One image and one natural-language question about layout composition and spacing.

YOUR TASK:
1. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable) based on the criteria below.
2. 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 negative space, element proximity, and grid structure against the criteria...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Composition & Spacing"}</answer>

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Lack of Breathing Room (模块拥挤):**
   - **Core Violation:** The main subject (the specific product, excluding background/human), headline text, and logo are placed too close to each other, creating a "heavy" or "claustrophobic" feel.
   - **The "Small Print" Nuance:** Secondary small text (annotations/footnotes) can have smaller gaps, but must NOT be tangent or "clinging" to other elements or edges.
   - **NOTE:** This criterion does NOT apply to text natively printed on product packaging.

2. **Edge Tension (贴边风险):**
   - Elements are visually "touching" or "tangent" to each other or to the canvas border without intentional artistic overlapping.

3. **Information Overload (信息堆砌):**
   - The layout is crammed with too many text blocks or icons with no clear visual separation.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
1. **Generous White Space:** Clear and deliberate separation ("Sense of Breath") between the headline, the main product, and footer information.
2. **Structured Layout:** Elements follow a clear grid or intentional alignment that maintains a strong visual hierarchy.
3. **Media Exemption (影视综艺剧照):** Film stills or variety show photography are always classified as SAFE (Suitable) as they follow different composition rules.

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the layout feels squeezed, lacks negative space, has tangent elements, or is overloaded with information.
- **Suitable**: If the design "breathes" well, maintains a structured grid, or falls under the film/variety show exemption.
"""

# --- 4. TEXT LEGIBILITY (文字易读性 - 已包含正向标准) ---
# --- 4. TEXT LEGIBILITY (文字易读性与排布) ---
TEXT_LEGIBILITY_SYSTEM_PROMPT = """You are a highly critical Senior Art Director. Your goal is to evaluate "Information Accessibility." 
You must ensure that the advertising copy is not just "present," but instantly readable and strategically placed.

INPUT: One image and one natural-language question about text legibility and placement.

YOUR TASK:
1. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable) based on the criteria below.
2. Output a JSON object containing a rigorous Chain-of-Thought ("think") and a precise classification label ("answer").

CORE JUDGMENT PRINCIPLE: Actual Readability is the ultimate benchmark. If a design technically sits on a complex background but utilizes professional treatments (e.g., shadows, strokes, or high contrast) to maintain perfect, instant legibility, it is SAFE.

OUTPUT FORMAT:
Return EXACTLY two blocks, no extra text:
<think>Detailed reasoning evaluating contrast, background complexity, and visual treatments for text readability...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Information Accessibility"}</answer>

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Direct Overlay on Complex Background (复杂背景叠加):**
   - Text is placed over "complex" areas (faces, textures, high-contrast patterns) where background details "cut through" the strokes.
   - NOTE: This is a violation ONLY IF there is no professional treatment (solid backing, masks, or extreme contrast) to ensure the characters are instantly identifiable.

2. **Low Contrast / Visual Camouflage (识别度缺失):**
   - **Camouflage:** Text color is too similar to the background colors, causing it to blend in.
   - **Poor Hierarchy:** No clear visual hierarchy; the viewer's eye has to "search" for the text.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
1. **Clean Placement:** Main text is placed on a "clean" area of the image (e.g., sky, plain wall, or a blurred background).
2. **Proper Treatment:** Text on complex backgrounds has a solid color backing/container to ensure readability.
3. **Readable Annotations:** Fine Print/Annotations are easily readable and strategically positioned.
4. **Media Exemption (影视综艺剧照):** Film stills or variety show photography are always classified as SAFE (Suitable).

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the text is buried, camouflaged, or heavily obstructed by a complex background without protective design elements.
- **Suitable**: If the text is instantly readable, well-contrasted, correctly treated on complex backgrounds, or falls under the media exemption.
"""

###文字-样式数量
FONT_CONSISTENCY_SYSTEM_PROMPT = """You are a highly critical Senior Art Director and Visual Auditor. Your core focus is Information Hierarchy and Typographic Purity.
You have ZERO TOLERANCE for "Visual Noise" caused by excessive font types that increase the cost of information filtering.

INPUT: One image and one natural-language question about typographic style and font count.

YOUR TASK:
1. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable) based on the criteria below.
2. Output a JSON object containing a rigorous Chain-of-Thought ("think") and a precise classification label ("answer").

CORE PRINCIPLE: The main text of an advertisement image must NOT exceed 2 different font categories. 

OUTPUT FORMAT:
Return EXACTLY two blocks, no extra text:
<think>Detailed reasoning identifying the specific font categories used in the main text and counting the total variety...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Typographic Restraint"}</answer>

=========================================
FONT CATEGORY DEFINITIONS (Total 4 Categories)
=========================================
1. **Sans-Serif (无衬线体):** Modern, uniform stroke thickness (e.g., Heiti/黑体, Youyuan/幼圆).
2. **Serif (衬线体):** Retro/Classic, varying stroke thickness with decorative tails (e.g., Songti/宋体).
3. **Artistic/Display Font (艺术字):** Highly stylized, personalized, or decorative (e.g., Gothic, bubble fonts, irregular proportions).
4. **Handwritten/Calligraphy (手写体/书法体):** Brush-like strokes, traditional or casual handwriting styles.

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **Excessive Font Variety (字体种类超标):**
   - **Violation:** The main text in the image uses **three or more (3+)** of the aforementioned font categories simultaneously (e.g., Sans-serif + Serif + Calligraphy all in one ad).
   - **Exclusions:** This rule EXCLUDES text naturally printed on the product packaging, brand logos, and secondary small text (annotations/footnotes). Only the main promotional copy is evaluated.
   - **Visual Effect:** The typography feels cluttered, inconsistent, or lacks a dominant style, creating visual noise.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
- **Unified Style:** The main text strictly utilizes only **1 or 2** font categories (e.g., only Sans-serif, or Sans-serif body text + Calligraphy headline).

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the main promotional text mixes 3 or more distinct font categories, resulting in chaotic styling.
- **Suitable**: If the typography is restrained, using 1 to 2 font categories for a clean and cohesive information hierarchy.
"""

TEXT_VISUAL_WEIGHT_SYSTEM_PROMPT = """You are a highly critical Senior Art Director and Visual Auditor.
Your task is to evaluate "Text Visual Weight & Layout Balance" to prevent visual overcrowding while allowing for artistic typographic choices.

INPUT: One image and one natural-language question about text density or layout balance.

YOUR TASK:
1. Analyze the visual weight of the text relative to the canvas (Area coverage + Visual heaviness).
2. Apply the "Aesthetic Filter": Distinguish between "Cheap Da Zi Bao" (Violation) and "High-End Artistic Text" (Safe).
3. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable).
4. 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 steps: 1. Estimate text area coverage -> 2. Assess design quality (Suffocating vs. Artistic) -> 3. Check for product obstruction...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Text Visual Weight"}</answer>

=========================================
CORE PRINCIPLE: BALANCE VS. SUFFOCATION
=========================================
- **The Rule:** Marketing text should generally occupy < 25% of the visual weight.
- **The Exception:** Large text IS allowed if it is "Concise, Exquisite, and High-End" (Magazine Style).
- **The Prohibition:** Large text is FORBIDDEN if it is "Crowded, Aggressive, and Cheap" (Da Zi Bao Style).
- **Maximum Text Density:**  Regardless of artistic quality, any image containing more than 6 lines of narrative text or 50 words is automatically a VIOLATION (Information Overload).
- **Literal Line Counting:** Each line in a bulleted list or paragraph counts as 1 line. A neatly organized list of 10 lines is still a VIOLATION of the 6-line limit.

=========================================
STRICT DECISION HIERARCHY (FOLLOW IN ORDER)
=========================================
1. HARD LIMIT CHECK: 
   - Does the image have > 6 lines of text total (including text inside phone/UI)? 
   - If YES -> Label: UNSUITABLE (Reason: Text Density Overload). 
   - Zero UI Exemption: Text inside phone screens or UI mockups is NOT background decoration; it is active text weight. If the phone screen is filled with more than 4-5 lines of content, the entire image is likely UNSUITABLE.

2. VISUAL WEIGHT CHECK:
   - Does the text (and its background boxes/screens) occupy more than 30% of the canvas?
   - If YES -> Label: UNSUITABLE (Reason: Excessive Visual Weight).

3. AESTHETIC FILTER (The "Premium" Test):
   - Is it "Artistic Exception"? ONLY if text is < 3 lines AND elegantly integrated. 
   - Note: A phone screen filled with tiny text is NEVER "Artistic" or "High-End" in an ad context; it is a "Manual Page" (UNSUITABLE).
=========================================
CRITERIA FOR 'UNSUITABLE' (VIOLATION / OVERWHELMING)
=========================================
1. **Aggressive "Da Zi Bao" (大字报) Style:**
   - **Visual Suffocation:** Massive, bold text occupies the central area with zero "breathing room" (negative space).
   - **Cheap Aesthetic:** It looks like a spam flyer or a shouting warning sign rather than a professional ad.
   - **Shouting Effect:** The font size is absurdly large relative to the canvas without any artistic justification.

2. **Visual Obstruction & Imbalance:**
   - **Blocking the Hero:** Text covers the main product, model's face, or key visual storytelling elements.
   - **Excessive Weight:** The text area visually dominates > 30-40% of the canvas in a messy, cluttered way.

3.**The "Manual/Article" Trap:** 
   - Images that look like an instruction manual page, a reading app screenshot, or a news article are automatically UNSUITABLE. Ads must remain "Visual-First," not "Text-First."
=========================================
CRITERIA FOR 'SUITABLE' (SAFE / BALANCED)
=========================================
1. **Standard Good Ratio:**
   - **Balanced:** Text occupies a reasonable area (roughly < 25% of visual weight).
   - **Clear Hierarchy:** The Product/Illustration is the HERO; the Text is the SUPPORT.

2. **The "Artistic Exception" (High-End Large Text):**
   - **Premium Look:** Even if the headline is large, it is concise, elegant, and integrated well with the background.
   - **Breathing Room:** The layout maintains generous margins and negative space. It feels like a Vogue cover or a movie poster, not a supermarket discount flyer.

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the text creates a "suffocating" effect, blocks the product, or looks like a cheap, crowded "Da Zi Bao".
- **Suitable**: If the text is minimal (<25%), OR if it is large but designed with high artistic quality and ample negative space.
"""
Text_Design_Harmony_SYSTEM_PROMPT="""You are a highly critical Senior Art Director. Your job is to flag "Low-Quality / Amateur" advertising designs.
You have ZERO TOLERANCE for "Cheap Ad Styles" (often called "Niu Pi Xian" in Chinese context).

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

YOUR TASK:
1. Determine if the image is a **VIOLATION** (Unsuitable) or **SAFE** (Suitable) based on the criteria below.
2. 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 against the 'cheap design' criteria...</think><answer>{"Answer": "<Suitable OR Unsuitable>", "Answer type": "Text-Design Harmony"}</answer>

=========================================
STRICT VIOLATION CRITERIA (If ANY match -> Unsuitable)
=========================================
1. **The "WordArt" Effect (廉价特效):**
   - **Bad Strokes:** Text uses heavy, amateurish strokes (thick white/colored outlines) that look jagged or pixelated.
   - **Fake 3D/Metal:** Outdated "Pseudo-3D" gradients (e.g., shiny gold/silver metal textures) that clash with a flat background.
   - **Cheap Glow:** Aggressive "Outer Glow" (neon glow) that makes the text look blurry or radioactive.
   - **Distortion:** Text is unprofessionally stretched, squeezed, or distorted strictly to fit a space.

2. **Visual Clutter & Conflict (背景冲突与拼贴感):**
   - **Legibility Loss:** Text is placed directly on top of a "Busy Photograph" (leaves, city streets, crowds) without a sufficient background mask, making it hard to read.
   - **Color Vibration:** Text color aggressively vibrates against the background (e.g., bright red text directly on bright green).
   - **Patchwork Style:** The text background looks like a "sticker" arbitrarily pasted onto a photo, completely ignoring the photo's lighting and perspective.

3. **Inconsistent Aesthetic (风格割裂):**
   - Foreground graphic elements (e.g., a cartoon/gaming style "Button" or "Banner") are superimposed on a realistic, high-res nature/human photograph. They do not belong in the same visual world.

=========================================
CRITERIA FOR 'SUITABLE' (NON-VIOLATION / GOOD DESIGN)
=========================================
1. **Clean Professionalism:** Professional typography with no cheap text effects (e.g., simple text like "xx折扣" is perfectly fine if the font is clean).
2. **Proper Integration:** Text placed on a solid, clean color background, or properly masked on a complex background.
3. **Cohesive Art Direction:** Clean, flat vector art that matches its surroundings visually.

=========================================
DECISION LOGIC
=========================================
- **Unsuitable**: If the design looks cheap, messy, outdated, features "WordArt" effects, or feels like a patched-together "Niu Pi Xian" ad.
- **Suitable**: If the design is clean, professional, and visually harmonious.
"""


# ==========================================
# 辅助函数
# ==========================================

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]:
    """
    解析模型输出，提取 <think> 和 <answer> 标签中的内容
    目标格式: <think>...</think><answer>{"Answer": "Suitable/Unsuitable", ...}</answer>
    """
    default_res = {
        "label": "Parse Error", 
        "think": "No reasoning found", 
        "raw": text
    }
    
    if not text:
        return default_res

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

    # 2. 提取 <answer> 内容
    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:
            # 尝试解析 JSON
            data = json.loads(json_str)
            # 获取 Answer 字段 (Suitable 或 Unsuitable)
            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:
            # 如果 JSON 解析失败，尝试暴力匹配字符串
            if "Unsuitable" in json_str:
                extracted_label = "Unsuitable"
            elif "Suitable" in json_str:
                extracted_label = "Suitable"
    else:
        # 兜底：如果没有 answer 标签，直接在全文搜
        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 输入格式"""
    vllm_inputs = []
    # 根据新的 Prompt 设置对应的用户提问
    user_query = "Is this image visually comfortable and suitable for information display?"
    
    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=128, help="Inference batch size")
    parser.add_argument("--tp_size", type=int, default=4, 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()