import re
import json
import os
from collections import defaultdict

def extract_data_from_log(log_content, extract_id=None):
    # 提取查询ID
    query_id_match = re.search(r'找到查询ID: (case_\w+_\d+)', log_content)
    query_id = query_id_match.group(1) if query_id_match else "unknown_id"
    
    # 尝试提取epoch信息 (如果存在)
    epoch_match = re.search(r'epoch[_-](\d+)', log_content, re.IGNORECASE)
    epoch = epoch_match.group(1) if epoch_match else None
    
    # 提取所有回复
    responses = []
    
    # 用正则表达式匹配每个模型回复块
    response_blocks = re.findall(
        r'模型回复 #(\d+) \(评分: ([\d.]+)\):\s*-+\s*<reasoning>(.*?)</reasoning>\s*<answer>(.*?)</answer>\s*-+\s*参考答案:\s*-+\s*(.*?)\s*=+',
        log_content, re.DOTALL
    )
    
    for idx, score, reasoning, answer, reference in response_blocks:
        # 将reasoning和answer合并为completion
        completion = f"<reasoning>{reasoning.strip()}</reasoning>\n<answer>{answer.strip()}</answer>"
        
        responses.append({
            "response_id": int(idx),
            "score": float(score),
            "completion": completion,
            "reference": reference.strip()
        })
    
    result = {
        "query_id": query_id,
        "responses": responses
    }
    
    # 添加提取ID和epoch信息(如果有)
    if extract_id:
        result["extract_id"] = extract_id
    if epoch:
        result["epoch"] = int(epoch)
        
    return result

def process_log_files(log_dir, output_file):
    all_data = []
    extract_counter = 1  # 用于生成唯一的extract_id
    
    # 如果输入是单个文件而不是目录
    if os.path.isfile(log_dir):
        with open(log_dir, 'r', encoding='utf-8') as f:
            log_content = f.read()
        
        # 生成唯一ID: 文件名_1
        base_filename = os.path.basename(log_dir).split('.')[0]
        extract_id = f"{base_filename}_{extract_counter}"
        
        data = extract_data_from_log(log_content, extract_id)
        all_data.append(data)
    else:
        # 处理目录中的所有文本文件
        for filename in sorted(os.listdir(log_dir)):
            if filename.endswith('.txt') or filename.endswith('.log'):
                file_path = os.path.join(log_dir, filename)
                
                with open(file_path, 'r', encoding='utf-8') as f:
                    log_content = f.read()
                
                # 生成唯一ID: 文件名_计数器
                base_filename = filename.split('.')[0]
                extract_id = f"{base_filename}_{extract_counter}"
                extract_counter += 1
                
                data = extract_data_from_log(log_content, extract_id)
                all_data.append(data)
    
    # 写入JSON文件
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(all_data, f, ensure_ascii=False, indent=2)
    
    print(f"已成功提取 {len(all_data)} 个问题的回复，并保存到 {output_file}")

# 使用示例
if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 3:
        print("用法: python extract_responses.py <日志文件或目录> <输出的JSON文件>")
        sys.exit(1)
    
    log_path = sys.argv[1]
    output_file = sys.argv[2]
    
    process_log_files(log_path, output_file)