import concurrent.futures
import multiprocessing
import threading
import time
import sys
import argparse
import json

class TestRunner:
    def __init__(self, max_workers=10):
        self.max_workers = max_workers
    
    def compile_code(self, code_str, main_function_name=None):
        try:
            local_vars = {}
            exec(code_str, local_vars)
            if main_function_name is not None:
                func = local_vars.get(main_function_name)
                return func if callable(func) else None
            return next((obj for obj in local_vars.values() if callable(obj)), None)
        except Exception as e:
            print(f"Compilation Error: {str(e)}, code_str:\n {code_str}")
            return None

    def _run_all_tests(self, functions, test_cases, timeout=None):
        tasks = []
        for func_name in functions:
            func_info = functions[func_name]
            for test_name in test_cases:
                tasks.append((
                    func_name,
                    test_name,
                    func_info['code'],
                    func_info['main_function_name'],
                    test_cases[test_name]['test_function']
                ))
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = []
            for task in tasks:
                futures.append(executor.submit(
                    self._run_single_test,
                    task[2],  # func_code
                    task[3],  # main_func_name
                    task[4],  # test_code
                    timeout
                ))
            
            function_results = {fn: {} for fn in functions}
            test_results = {tn: {} for tn in test_cases}
            
            for (func_name, test_name, *_), future in zip(tasks, futures):
                result = future.result()
                function_results[func_name][test_name] = result
                test_results[test_name][func_name] = result
        
        return function_results, test_results
    
    def _run_single_test(self, func_code, main_func_name, test_code, timeout):
        manager = multiprocessing.Manager()
        queue = manager.Queue()
        process = multiprocessing.Process(
            target=self._worker_process,
            args=(func_code, main_func_name, test_code, timeout, queue)
        )
        process.start()
        process.join(timeout + 1)  # 稍等片刻确保进程清理
        
        if process.is_alive():
            process.terminate()
            process.join()
            return False
        
        return queue.get() if not queue.empty() else False
    
    def _worker_process(self, func_code, main_func_name, test_code, timeout, queue):
        try:
            # 动态加载被测试函数
            func_namespace = {}
            exec(func_code, func_namespace)

            func = self.compile_code(func_code, main_function_name=main_func_name)
            
            # 动态加载测试函数
            test_namespace = {}
            exec(test_code, test_namespace)
            test_func = self.compile_code(test_code, main_function_name=None)
            
            # 使用线程控制执行时间
            result_container = []
            event = threading.Event()
            
            def worker_thread():
                try:
                    result_container.append(test_func(func))
                except Exception as e:
                    result_container.append(False)
                finally:
                    event.set()
            
            thread = threading.Thread(target=worker_thread)
            thread.daemon = True
            thread.start()
            
            event.wait(timeout)
            
            if not event.is_set() or not result_container:
                queue.put(False)
            else:
                queue.put(bool(result_container[0]))
        except Exception as e:
            queue.put(False)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--functions", type=str, required=True, help="JSON string for functions dictionary")
    parser.add_argument("--test_cases", type=str, required=True, help="JSON string for test cases dictionary")
    parser.add_argument("--max_workers", type=int, default=10, help="Maximum number of workers")
    parser.add_argument("--timeout", type=float, default=5, help="Timeout value in seconds")
    
    args = parser.parse_args()
    try:
        functions = json.loads(args.functions)
        test_cases = json.loads(args.test_cases)
    except Exception as e:
        print("Error parsing JSON input:", e)
        sys.exit(1)
    
    runner = TestRunner(max_workers=args.max_workers)
    start_time = time.time()
    function_results, test_results = runner._run_all_tests(functions, test_cases, timeout=args.timeout)
    elapsed_time = time.time() - start_time
    
    print("Function Results:", function_results)
    print("Test Results:", test_results)
    print("Execution Time:", elapsed_time)

if __name__=="__main__":
    main()
