import os
import json
import subprocess
import pandas as pd
from pathlib import Path

GLOBAL_MODEL_PATH = "../stage2_object_v8_1200"

##示例:每个都传入绝对路径
TASK_CONFIGS = [
    ("rule_11_vllm.py", "文字占比"),
    ("rule_12_vllm.py", "文字-样式数量合理"),
    #("rule_13_vllm.py", "文字-排布位置"),
    #("rule_14_vllm.py", "文字-设计搭配协调"),
    ("rule_17_vllm.py", "信息量"),
    ("rule_18_vllm.py", "排布间距"),
    ("rule_19_vllm.py", "内容构图"),
]

SUMMARY_OUTPUT = "vllm_audit_summary.csv"

def run_vllm_task(script, input_dir):

    cmd = [
        "python", script,
        "--input_dir", input_dir,
        "--model_path", GLOBAL_MODEL_PATH
    ]
    print(f"\n[EXEC] 正在运行: {script}")
    print(f"[PATH] 输入目录: {Path(input_dir).name}")
    
    try:
        subprocess.run(cmd, check=True)
        return True
    except subprocess.CalledProcessError as e:
        print(f"[ERROR] 脚本 {script} 运行失败: {e}")
        return False

def extract_metrics(script, input_dir):
    input_path = Path(input_dir)
    json_file = input_path / f"audit_result_{input_path.name}.json"
    
    if not json_file.exists():
        return {
            "规则名称": script, 
            "对应目录": input_path.name, 
            "结果": "未找到 JSON 文件"
        }

    with open(json_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    total = len(data)
    unsuitable = sum(1 for item in data if item["label"] == "Unsuitable")
    suitable = sum(1 for item in data if item["label"] == "Suitable")
    recall = (unsuitable / total) if total > 0 else 0

    return {
        "规则名称": script.replace("_vllm.py", ""),
        "对应目录": input_path.name,
        "样本总数": total,
        "不通过数(Unsuitable)": unsuitable,
        "通过数(Suitable)": suitable,
        "召回率(违规检出率)": f"{recall:.2%}",
        "所用模型": Path(GLOBAL_MODEL_PATH).name
    }

def main():
    final_results = []
    
    for script, input_dir in TASK_CONFIGS:
        success = run_vllm_task(script, input_dir)
        
        metrics = extract_metrics(script, input_dir)
        final_results.append(metrics)
        
    df = pd.DataFrame(final_results)
    df.to_csv(SUMMARY_OUTPUT, index=False, encoding='utf-8-sig')
    
    print("\n" + "="*70)
    print(f"所有任务运行完毕！汇总报告已保存至: {SUMMARY_OUTPUT}")
    print("-" * 70)
    print(df.to_string(index=False))
    print("="*70)

if __name__ == "__main__":
    main()