#!/usr/bin/env python3
"""
信息图表质量检查脚本

使用Gemini VLLM模型检查图表质量，自动筛选合格的图表并复制到目标文件夹。

使用方法:
    python check_charts.py <source_folder> <target_folder>
"""

import json
import base64
import shutil
import argparse
import os
from pathlib import Path
from typing import Dict, Optional, Tuple
from multiprocessing import Pool, Manager
from google import genai
from google.genai import types

# 尝试导入配置
try:
    from config import LLM_API_KEY, LLM_BASE_URL, LLM_MODEL
except ImportError:
    print("⚠️  警告: config.py 未找到，请设置环境变量或修改下面的默认值")
    LLM_API_KEY = os.getenv("GEMINI_API_KEY") or os.getenv("OPENAI_API_KEY", "")
    LLM_BASE_URL = os.getenv("GEMINI_BASE_URL", "https://aihubmix.com/gemini")
    LLM_MODEL = os.getenv("GEMINI_MODEL", "gemini-3-flash-preview")


class ChartQualityChecker:
    """使用Gemini VLLM检查信息图表质量"""
    
    def __init__(self, api_key=None, base_url=None, model=None):
        """
        初始化检查器
        
        Args:
            api_key: API密钥
            base_url: API基础URL
            model: 模型名称
        """
        self.api_key = api_key or LLM_API_KEY
        self.base_url = base_url or LLM_BASE_URL or 'https://aihubmix.com/gemini'
        self.model = model or LLM_MODEL or 'gemini-3-flash-preview'
        
        # 初始化Google genai client
        self.client = genai.Client(
            api_key=self.api_key,
            http_options={'base_url': self.base_url}
        )
    
    def check_chart(self, image_path: str) -> Tuple[bool, str]:
        """
        检查单个图表是否符合质量标准
        
        Args:
            image_path: 图表图片路径
            
        Returns:
            (是否通过, 原因说明)
        """
        try:
            # 编码图片
            image_base64 = self._encode_image(image_path)
            if not image_base64:
                return False, "图片编码失败"
            
            # 构建检查提示
            prompt = """Please evaluate this image as a data visualization or infographic. Check whether:

1. It contains at least 4 meaningful data points. 数据不全为0.
2. The information is complete (no missing labels, legends, units, or explanations).
3. There is no text obstruction (no overlapping, cropped, blurred, or unreadable text).
4. The design is reasonable and clear, including proper layout, visual hierarchy, and readability.

Return a JSON object with this EXACT format:
{
  "passed": true/false,
  "reason": "Brief explanation in Chinese",
  "details": {
    "data_points": "pass/fail - explanation",
    "completeness": "pass/fail - explanation",
    "text_quality": "pass/fail - explanation",
    "design_clarity": "pass/fail - explanation"
  }
}

IMPORTANT: All checks must pass for overall "passed" to be true.
"""
            
            # 查询VLLM
            result = self._query_vllm(prompt, image_base64)
            
            if not result:
                return False, "VLLM查询失败"
            
            passed = result.get('passed', False)
            reason = result.get('reason', '未知原因')
            
            return passed, reason
            
        except Exception as e:
            return False, f"检查过程出错: {str(e)}"
    
    def _encode_image(self, image_path: str) -> Optional[str]:
        """编码图片为base64"""
        try:
            path = Path(image_path)
            if not path.exists():
                return None
            
            with open(path, 'rb') as f:
                image_data = f.read()
            
            return base64.b64encode(image_data).decode('utf-8')
        except Exception as e:
            print(f"❌ 图片编码错误 {image_path}: {e}")
            return None
    
    def _query_vllm(self, prompt: str, image_base64: str) -> Optional[Dict]:
        """
        查询VLLM API
        
        Args:
            prompt: 文本提示
            image_base64: base64编码的图片
            
        Returns:
            解析后的JSON响应
        """
        try:
            # 构建multimodal内容
            contents = [
                types.Part(text="You are an expert in data visualization quality assessment. Always return valid JSON only."),
                types.Part(text=prompt),
                types.Part(inline_data=types.Blob(
                    mime_type="image/png",
                    data=image_base64
                ))
            ]
            
            # 调用genai API，关闭thinking模式
            response = self.client.models.generate_content(
                model=self.model,
                contents=contents,
                config=types.GenerateContentConfig(
                    temperature=0.2
                )
            )
            
            content = response.text.strip()
            
            # 清理可能的markdown代码块
            if content.startswith('```'):
                lines = content.split('\n')
                content = '\n'.join(lines[1:-1] if lines[-1].strip() == '```' else lines[1:])
                content = content.replace('```json', '').replace('```', '').strip()
            
            # 解析JSON
            return json.loads(content)
            
        except json.JSONDecodeError as e:
            print(f"❌ VLLM响应不是有效的JSON: {e}")
            return None
        except Exception as e:
            print(f"❌ VLLM查询错误: {e}")
            return None


def process_single_folder(args: Tuple) -> Dict:
    """
    处理单个子文件夹（用于多进程）
    
    Args:
        args: (folder_path, target_folder, api_key, base_url, model, record_file, lock)
        
    Returns:
        处理结果字典
    """
    folder_path, target_folder, api_key, base_url, model, record_file, lock = args
    
    folder = Path(folder_path)
    chart_png = folder / "chart.png"
    folder_name = folder.name
    
    result = {
        'folder': str(folder),
        'folder_name': folder_name,
        'success': False,
        'reason': '',
        'copied': False
    }
    
    # 检查chart.png是否存在
    if not chart_png.exists():
        result['reason'] = 'chart.png不存在'
        return result
    
    try:
        # 创建检查器实例（每个进程独立）
        checker = ChartQualityChecker(api_key=api_key, base_url=base_url, model=model)
        
        # 检查图表质量
        passed, reason = checker.check_chart(str(chart_png))
        
        result['success'] = passed
        result['reason'] = reason
        
        if passed:
            # 复制文件到目标文件夹
            target_subfolder = Path(target_folder) / folder.name
            target_subfolder.mkdir(parents=True, exist_ok=True)
            
            # 需要复制的文件列表
            files_to_copy = ['chart.png', 'chart.svg', 'data.json', 'info.json']
            
            copied_files = []
            for filename in files_to_copy:
                source_file = folder / filename
                if source_file.exists():
                    target_file = target_subfolder / filename
                    shutil.copy2(source_file, target_file)
                    copied_files.append(filename)
            
            result['copied'] = True
            result['copied_files'] = copied_files
            
            # 记录已检查的文件夹（使用锁确保多进程安全）
            if record_file:
                with lock:
                    with open(record_file, 'a', encoding='utf-8') as f:
                        f.write(f"{folder_name}\n")
        
        return result
        
    except Exception as e:
        result['reason'] = f'处理错误: {str(e)}'
        return result


def load_checked_folders(record_file: str) -> set:
    """
    加载已检查过的文件夹列表
    
    Args:
        record_file: 记录文件路径
        
    Returns:
        已检查文件夹名称的集合
    """
    record_path = Path(record_file)
    if not record_path.exists():
        return set()
    
    try:
        with open(record_path, 'r', encoding='utf-8') as f:
            return set(line.strip() for line in f if line.strip())
    except Exception as e:
        print(f"⚠️  读取记录文件失败: {e}")
        return set()


def check_charts_parallel(source_folder: str, target_folder: str, num_processes: int = 10, record_file: str = "record.txt"):
    """
    并行检查文件夹下的所有图表
    
    Args:
        source_folder: 源文件夹路径
        target_folder: 目标文件夹路径
        num_processes: 并行进程数
        record_file: 记录文件路径
    """
    source_path = Path(source_folder)
    target_path = Path(target_folder)
    
    if not source_path.exists():
        print(f"❌ 源文件夹不存在: {source_folder}")
        return
    
    # 创建目标文件夹
    target_path.mkdir(parents=True, exist_ok=True)
    
    # 加载已检查过的文件夹
    print(f"📋 加载检查记录: {record_file}")
    checked_folders = load_checked_folders(record_file)
    print(f"   已有 {len(checked_folders)} 个文件夹的检查记录")
    
    # 查找所有包含chart.png的子文件夹
    print(f"🔍 扫描文件夹: {source_folder}")
    subfolders = []
    skipped_count = 0
    
    for item in source_path.iterdir():
        if item.is_dir():
            chart_png = item / "chart.png"
            if chart_png.exists():
                # 检查是否已经检查过
                if item.name in checked_folders:
                    skipped_count += 1
                else:
                    subfolders.append(item)
    
    print(f"📊 找到 {len(subfolders) + skipped_count} 个包含chart.png的子文件夹")
    print(f"   ⏭️  跳过 {skipped_count} 个已检查的文件夹")
    print(f"   🆕 需要检查 {len(subfolders)} 个新文件夹")
    
    if not subfolders:
        print("⚠️  没有需要检查的新图表")
        return
    
    # 准备多进程参数（添加Manager用于进程间共享锁）
    manager = Manager()
    lock = manager.Lock()
    
    args_list = [
        (str(folder), str(target_path), LLM_API_KEY, LLM_BASE_URL, LLM_MODEL, record_file, lock)
        for folder in subfolders
    ]
    
    # 使用进程池并行处理
    print(f"🚀 开始并行检查 (使用 {num_processes} 个进程)...\n")
    
    passed_count = 0
    failed_count = 0
    
    with Pool(processes=num_processes) as pool:
        # 使用imap_unordered以便及时显示结果
        results = pool.imap_unordered(process_single_folder, args_list)
        
        for i, result in enumerate(results, 1):
            folder_name = result['folder_name']
            
            if result['success']:
                passed_count += 1
                status = "✅ 通过"
                copied_info = f" | 已复制 {len(result.get('copied_files', []))} 个文件"
            else:
                failed_count += 1
                status = "❌ 未通过"
                copied_info = ""
                
                # 即使未通过也记录，避免重复检查
                with lock:
                    with open(record_file, 'a', encoding='utf-8') as f:
                        f.write(f"{folder_name}\n")
            
            print(f"[{i}/{len(subfolders)}] {status} | {folder_name}")
            print(f"    原因: {result['reason']}{copied_info}")
    
    # 统计结果
    print(f"\n{'='*60}")
    print(f"📈 检查完成!")
    print(f"{'='*60}")
    print(f"  本次检查: {len(subfolders)} 个图表")
    print(f"  ✅ 通过: {passed_count} 个 ({passed_count/len(subfolders)*100:.1f}%)")
    print(f"  ❌ 未通过: {failed_count} 个 ({failed_count/len(subfolders)*100:.1f}%)")
    print(f"  ⏭️  已跳过: {skipped_count} 个（之前已检查）")
    print(f"  📁 目标文件夹: {target_folder}")
    print(f"  📋 记录文件: {record_file}")
    print(f"{'='*60}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='检查信息图表质量并复制合格的图表',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
    python check_charts.py ./output ./checked_output
    python check_charts.py /path/to/source /path/to/target --processes 5
    python check_charts.py ./output ./checked_output --record my_record.txt
        """
    )
    
    parser.add_argument('source', help='源文件夹路径（包含多个子文件夹，每个子文件夹有chart.png）')
    parser.add_argument('target', help='目标文件夹路径（存放通过检查的图表）')
    parser.add_argument('--processes', '-p', type=int, default=20,
                        help='并行进程数 (默认: 10)')
    parser.add_argument('--record', '-r', type=str, default='record.txt',
                        help='检查记录文件路径 (默认: record.txt)')
    
    args = parser.parse_args()
    
    # 检查API密钥
    if not LLM_API_KEY:
        print("❌ 错误: 未配置LLM_API_KEY")
        print("请在config.py中配置API密钥或设置环境变量")
        return
    
    # 执行检查
    check_charts_parallel(args.source, args.target, args.processes, args.record)


if __name__ == '__main__':
    main()