"""
This code is modified from the evaluation code of MathInstruct.
https://github.com/TIGER-AI-Lab/MAmmoTH/blob/main/math_eval/code_utils.py
"""

import multiprocessing
import threading
from contextlib import redirect_stdout
from io import StringIO


def format_code(code_str: str):
    code = "def run_it():\n"
    for line in code_str.split("\n"):
        code += "  " + line + "\n"
    code += "run_it()"
    return code


class CodeExecutor:
    def __init__(self, code: str, timeout: int, use_process: bool):
        self.code = format_code(code)
        self.timeout = timeout
        self.error = ""
        self.use_process = use_process

    def execute_code(self, return_val):
        try:
            f = StringIO()
            with redirect_stdout(f):
                exec(self.code, globals(), locals())
            s = f.getvalue()
            s = s.strip("\n")
            return_val["result"] = s
        except Exception:
            pass

    @staticmethod
    def execute_code_with_string(code, index, return_val):
        code = format_code(code)
        try:
            f = StringIO()
            with redirect_stdout(f):
                exec(code, globals(), locals())
            s = f.getvalue()
            s = s.strip("\n")
            return_val[index] = s
        except Exception as e:
            pass

    def run(self):
        if self.use_process:
            manager = multiprocessing.Manager()
            return_dict = manager.dict()
            process = multiprocessing.Process(target=self.execute_code, args=(return_dict,))
            process.start()
            process.join(timeout=self.timeout)
            process.terminate()
        else:
            return_dict = {}
            thread = threading.Thread(target=self.execute_code, args=(return_dict,))
            thread.start()
            thread.join(timeout=self.timeout)
            if thread.is_alive():
                thread.join()  # Ensures the thread is terminated before continuing
                print("time out!")
                self.error = "Execution timed out"

        if "result" in return_dict:
            return return_dict["result"]
        else:
            return ""
