import json
import os
import logging
from typing import List, Dict, Any, Tuple, Optional, Set
from collections import defaultdict
import tempfile
import subprocess
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
import re
import hashlib
import queue
import io
import traceback

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("testcase_execution.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 用于线程安全地写入文件的锁
file_lock = threading.Lock()

def extract_python_code_block(text: str) -> Optional[str]:
    """
    从文本中提取Python代码块 (```python ... ```)
    
    Args:
        text: 包含代码块的文本
        
    Returns:
        提取出的Python代码，如果没有找到则返回None
    """
    # 查找```python开头，```结尾的代码块
    pattern = r'```python\s*(.*?)\s*```'
    # 使用re.DOTALL标志使.能匹配换行符
    match = re.search(pattern, text, re.DOTALL)
    
    if match:
        return match.group(1).strip()
    
    # 如果没有找到python标记的代码块，尝试查找一般的代码块
    pattern = r'```\s*(.*?)\s*```'
    match = re.search(pattern, text, re.DOTALL)
    
    if match:
        return match.group(1).strip()
    
    # 如果仍然没有找到代码块，返回原始文本(可能本身就是代码)
    return text

def execute_test_code_in_subprocess(test_code: str, times: int = 10) -> List[Any]:
    """
    将测试代码写入临时Python文件并执行，捕获所有输出作为结果
    
    Args:
        test_code: 要执行的Python代码
        times: 执行次数
        
    Returns:
        长度为times的结果列表
    """
    results = []
    
    # 创建临时目录
    with tempfile.TemporaryDirectory() as temp_dir:
        # 将原始代码写入临时文件
        original_file_path = os.path.join(temp_dir, "original_code.py")
        with open(original_file_path, 'w') as f:
            f.write(test_code)
        
        # 创建一个包装文件，输出原始代码的全部输出
        wrapper_file_path = os.path.join(temp_dir, "wrapper.py")
        with open(wrapper_file_path, 'w') as f:
            f.write(f"""
import subprocess
import sys
import json

# 执行原始代码并捕获输出
result = subprocess.run(
    ["{sys.executable}", "{original_file_path}"],
    capture_output=True,
    text=True,
    timeout=30
)

# 输出所有非空行
if result.returncode == 0:
    output_lines = [line.strip() for line in result.stdout.split('\\n') if line.strip()]
    if output_lines:
        # 输出完整内容（多行）
        print(json.dumps(output_lines))
    else:
        print("[]")
else:
    print(f"ERROR: {{result.stderr.strip()}}")
""")
        
        # 执行指定次数
        for _ in range(times):
            try:
                # 执行包装代码
                process = subprocess.run(
                    [sys.executable, wrapper_file_path],
                    capture_output=True,
                    text=True,
                    timeout=30  # 包装代码有自己的30秒超时，这里设置更长时间
                )
                
                if process.returncode == 0:
                    # 成功执行，获取输出
                    output = process.stdout.strip()
                    if output:
                        try:
                            # 解析JSON数组形式的输出行
                            output_lines = json.loads(output)
                            if isinstance(output_lines, list):
                                # 无论行数多少，都将所有行用换行符连接成一个字符串
                                results.append('\n'.join(output_lines))
                            else:
                                results.append(output)
                        except json.JSONDecodeError:
                            results.append(output)
                    else:
                        results.append(None)
                else:
                    # 执行错误
                    error_msg = process.stderr.strip()
                    logger.error(f"代码执行错误: {error_msg}")
                    results.append(f"ERROR: {error_msg}")
                    
            except subprocess.TimeoutExpired:
                logger.error(f"代码执行超时")
                results.append("ERROR: 执行超时")
            except Exception as e:
                logger.error(f"执行异常: {str(e)}")
                results.append(f"ERROR: {str(e)}")
    
    return results

def generate_case_id(case: Dict, problem_id: str, sample_id: Any) -> str:
    """
    为测试用例生成唯一标识符
    
    Args:
        case: 测试用例
        problem_id: 问题ID
        sample_id: 样例ID
        
    Returns:
        唯一标识符字符串
    """
    # 使用问题ID、样例ID和测试代码的哈希值作为唯一标识符
    test_code = case.get("test_code", "")
    key_str = f"{problem_id}_{sample_id}_{test_code}"
    
    # 使用MD5计算哈希值
    hash_object = hashlib.md5(key_str.encode())
    return hash_object.hexdigest()

def process_case(case: Dict, problem_id: str, sample_id: Any, times_per_test: int) -> Tuple[Dict, bool, str]:
    """
    处理单个测试用例并返回结果
    
    Args:
        case: 单个测试用例
        problem_id: 问题ID
        sample_id: 样例ID
        times_per_test: 每个测试执行的次数
        
    Returns:
        更新后的测试用例、是否成功的标志和用例ID
    """
    case_id = generate_case_id(case, problem_id, sample_id)
    
    test_code = case.get("test_code")
    if not test_code:
        logger.warning(f"问题 {problem_id}, 样例 {sample_id} 没有测试代码")
        return None, False, case_id
    
    # 从文本中提取Python代码块
    code = extract_python_code_block(test_code)
    if not code:
        logger.warning(f"问题 {problem_id}, 样例 {sample_id} 无法提取Python代码块")
        return None, False, case_id
        
    # 执行提取的代码
    all_results = execute_test_code_in_subprocess(code, times_per_test)
    
    # 去重处理，保持原有顺序
    unique_results = []
    seen = set()
    for result in all_results:
        # 对于None和字符串类型可以直接去重
        result_key = str(result) if result is not None else None
        if result_key not in seen:
            seen.add(result_key)
            unique_results.append(result)
    
    # 更新结果
    updated_case = case.copy()
    updated_case["test_results"] = unique_results
    updated_case["case_id"] = case_id  # 添加case_id到结果中，方便后续判断
    
    # 检查结果JSON大小，如果超过50MB则跳过
    json_result = json.dumps(updated_case)
    if len(json_result.encode('utf-8')) > 50 * 1024 * 1024:  # 50MB
        logger.warning(f"问题 {problem_id}, 样例 {sample_id} 的结果超过50MB，已跳过")
        return None, False, case_id
    
    return updated_case, True, case_id

def write_result_to_file(output_file: str, result: Dict):
    """
    以线程安全的方式将结果写入文件
    
    Args:
        output_file: 输出文件路径
        result: 要写入的结果
    """
    with file_lock:
        with open(output_file, 'a') as out_f:
            out_f.write(json.dumps(result) + "\n")

def get_processed_case_ids(output_file: str) -> Set[str]:
    """
    从输出文件中获取已处理的测试用例ID
    
    Args:
        output_file: 输出文件路径
    
    Returns:
        已处理测试用例ID的集合
    """
    processed_ids = set()
    if os.path.exists(output_file):
        try:
            with open(output_file, 'r') as f:
                for line in f:
                    try:
                        case = json.loads(line.strip())
                        case_id = case.get("case_id")
                        if case_id:
                            processed_ids.add(case_id)
                    except json.JSONDecodeError:
                        continue
            logger.info(f"从输出文件中读取了 {len(processed_ids)} 个已处理的测试用例ID")
        except Exception as e:
            logger.error(f"读取输出文件失败: {e}")
    
    return processed_ids

class TaskQueueManager:
    """
    任务队列管理器，用于控制并发数量
    """
    def __init__(self, output_file: str, max_concurrent: int):
        self.queue = queue.Queue()
        self.output_file = output_file
        self.max_concurrent = max_concurrent
        self.active_tasks = 0
        self.task_lock = threading.Lock()
        self.total_completed = 0
        self.total_tasks = 0
        self.start_time = time.time()
        
    def add_task(self, task_fn, *args, **kwargs):
        """添加任务到队列"""
        self.total_tasks += 1
        self.queue.put((task_fn, args, kwargs))
    
    def worker(self):
        """工作线程，从队列获取任务并执行"""
        while True:
            try:
                # 从队列获取任务
                task_fn, args, kwargs = self.queue.get(block=False)
                
                # 增加活动任务计数
                with self.task_lock:
                    self.active_tasks += 1
                
                try:
                    # 执行任务
                    result = task_fn(*args, **kwargs)
                    
                    # 处理成功的结果
                    if result:
                        updated_case, success, case_id = result
                        if success and updated_case:
                            write_result_to_file(self.output_file, updated_case)
                            
                            with self.task_lock:
                                self.total_completed += 1
                                completed = self.total_completed
                                
                                # 每10个测试用例记录一次进度
                                if completed % 10 == 0:
                                    elapsed = time.time() - self.start_time
                                    remaining = self.total_tasks - completed
                                    logger.info(f"已处理 {completed}/{self.total_tasks} 个测试用例，用时 {elapsed:.2f} 秒，剩余 {remaining} 个")
                                
                except Exception as e:
                    logger.error(f"任务执行失败: {str(e)}")
                    logger.error(traceback.format_exc())
                
                finally:
                    # 减少活动任务计数
                    with self.task_lock:
                        self.active_tasks -= 1
                    
                    # 标记队列任务完成
                    self.queue.task_done()
                
            except queue.Empty:
                # 如果队列为空，线程退出
                break
    
    def run(self):
        """运行所有任务，控制并发数量"""
        if self.total_tasks == 0:
            logger.info("没有任务需要处理")
            return 0
        
        logger.info(f"开始处理 {self.total_tasks} 个任务，最大并发数: {self.max_concurrent}")
        
        # 创建工作线程池
        threads = []
        for _ in range(self.max_concurrent):
            thread = threading.Thread(target=self.worker)
            thread.daemon = True
            thread.start()
            threads.append(thread)
        
        # 等待所有任务完成
        self.queue.join()
        
        # 等待所有线程结束
        for thread in threads:
            thread.join(0.1)  # 短暂等待，不阻塞主线程
        
        elapsed = time.time() - self.start_time
        logger.info(f"任务处理完成，共处理 {self.total_completed}/{self.total_tasks} 个测试用例，总用时 {elapsed:.2f} 秒")
        
        return self.total_completed

def process_jsonl_file(input_file: str, output_file: str, times_per_test: int = 10, max_workers: int = None):
    """
    处理JSONL文件中的test_code，每个test_code执行指定次数，并将结果保存到输出文件
    支持断点续传，跳过已处理的测试用例，使用可控的并发数量
    
    Args:
        input_file: 输入的JSONL文件路径
        output_file: 输出的JSONL文件路径
        times_per_test: 每个测试执行的次数
        max_workers: 最大并行线程数，默认为CPU核心数x2
    """
    if not os.path.exists(input_file):
        logger.error(f"输入文件不存在: {input_file}")
        return
    
    # 如果未指定worker数量，使用CPU核心数的2倍（因为线程开销比进程小）
    if max_workers is None:
        import multiprocessing
        max_workers = multiprocessing.cpu_count() * 2
    
    logger.info(f"开始处理文件: {input_file}")
    logger.info(f"每个测试代码将执行 {times_per_test} 次")
    logger.info(f"最大并行线程数: {max_workers}")
    
    # 获取已处理的测试用例ID
    processed_ids = get_processed_case_ids(output_file)
    
    # 读取所有测试用例
    test_cases = []
    with open(input_file, 'r') as f:
        for line in f:
            try:
                test_cases.append(json.loads(line))
            except json.JSONDecodeError:
                logger.error(f"JSON解析错误: {line}")
    
    logger.info(f"读取了 {len(test_cases)} 个测试用例")
    
    # 按problem_id和sample_id组织测试用例
    problem_samples = defaultdict(list)
    for test_case in test_cases:
        problem_id = test_case.get("problem_id", "unknown")
        sample_id = test_case.get("sample_id", 0)
        problem_samples[(problem_id, sample_id)].append(test_case)
    
    logger.info(f"共有 {len(problem_samples)} 个不同的问题/样例组合")
    
    # 准备任务列表
    tasks = []
    for (problem_id, sample_id), cases in problem_samples.items():
        for case in cases:
            # 生成测试用例ID，如果已处理则跳过
            case_id = generate_case_id(case, problem_id, sample_id)
            if case_id in processed_ids:
                logger.info(f"跳过已处理的测试用例: 问题 {problem_id}, 样例 {sample_id}")
                continue
            tasks.append((case, problem_id, sample_id))
    
    logger.info(f"需要处理 {len(tasks)} 个新测试用例")
    
    # 如果没有新任务，直接返回
    if not tasks:
        logger.info("所有测试用例已处理完成，无需进一步操作")
        return
    
    # 确保输出文件存在
    if not os.path.exists(output_file):
        with open(output_file, 'w'):
            pass
    
    # 创建任务队列管理器
    task_manager = TaskQueueManager(output_file, max_workers)
    
    # 添加所有任务到队列
    for case, problem_id, sample_id in tasks:
        task_manager.add_task(process_case, case, problem_id, sample_id, times_per_test)
    
    # 运行任务队列管理器，执行所有任务
    total_processed = task_manager.run()
    
    logger.info(f"处理完成，本次共处理 {total_processed} 个新测试用例，总共 {len(processed_ids) + total_processed} 个测试用例，结果已保存到 {output_file}")

def main():
    """主函数"""
    try:
        # 设置参数
        input_file = "/home/superbench/xinzhang3/haoling/epicoder2/test_case_generator/question_testcode/V3_problems_V3_testcode_0422.jsonl"
        output_file = "/home/superbench/xinzhang3/haoling/epicoder2/test_case_generator/question_testcode_testcase/V3_problems_V3_testcode_0422.jsonl"
        times_per_test = 20  # 每个测试代码执行20次
        
        # 获取CPU核心数
        import multiprocessing
        cpu_count = multiprocessing.cpu_count()
        # 设置最大并行线程数
        max_workers = min(cpu_count * 2, 32)
        
        # 处理测试用例
        process_jsonl_file(input_file, output_file, times_per_test, max_workers)
        
    except KeyboardInterrupt:
        logger.info("程序被用户中断")
    except Exception as e:
        logger.error(f"程序执行出错: {str(e)}")
        logger.error(traceback.format_exc())

if __name__ == "__main__":
    main()
