import os
import subprocess
import json

# Set the directory where your .lean files are located

# Get a list of all .lean files in the directory
# lean_files = [f for f in os.listdir(directory) if f.endswith(".lean")]
# lean_files = ["test/file.lean"]
def main(args):
    command_list = []
    for i in range(8):
        with open(f"{args.input_path}/{i}.json", 'r', encoding='utf-8') as rf:
            for line in rf.readlines():
                try:
                    json_item = json.loads(line)
                    json_item['cmd']  = '\n'.join()
                except:
                    import pdb
                    pdb.set_trace()
                command_list.append(json_item)
    results = []
    passed = 0
    total = 0

    for item in command_list:
        data = '{"cmd": "%s", "allTactics": true}' % item['cmd']
        command = 'echo \'%s\' | lake exe repl' % data

        try:
            result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            stdout = result.stdout.decode('utf-8')
            stderr = result.stderr.decode('utf-8')
            results.append({
                'file_path': item['file_path'],
                'stdout': stdout,
                'stderr': stderr,
                'status': 'pass'
            })
            passed += 1
        except subprocess.CalledProcessError as e:
            results.append({
                'file_path': item['file_path'],
                'error': str(e),
                'status': 'nopass'
            })
        total += 1

    # Calculate pass rate
    pass_rate = passed / total * 100

    # Save results to a JSON file
    with open('results.json', 'w') as f:
        json.dump({'results': results, 'pass_rate': pass_rate}, f, indent=2)

if __name__ == '__main__':
    arg_parser = ArgumentParser()
    arg_parser.add_argument('--data_path', type=str,
                            default='data/grade-school-math-master/grade_school_math/data/test.jsonl')
    arg_parser.add_argument('--input_path', type=str, default='')
    arg_parser.add_argument('--cuda_num', type=int, default=8)
    arg_parser.add_argument('--output_path', type=str, default='total.json')
    arg_parser.add_argument('--generate_method', type=str,
                            choices=['single', 'sft', 'comp', 'self_consistency', 'single_consistency'])
    arg_parser.add_argument('--method', type=str, choices=['main', 'test', 'get_data'])
    args = arg_parser.parse_args()
    main(args)