import os
import io
import subprocess
import re
import pickle
import traceback
from contextlib import redirect_stdout
import tempfile

def run_task(code: str, timeout):
    running_code = "\n".join(code)
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as tmp_file:
        tmp_file.write(running_code)
        tmp_file_path = tmp_file.name

    try:
        result = subprocess.run(
            ["/root/anaconda3/bin/conda", "run", "-n", "darl-run", "python", tmp_file_path],
            capture_output=True,
            timeout=timeout
        )
        stdout = result.stdout.decode("utf-8")
        stderr = result.stderr.decode("utf-8")
        
        filtered_stderr = "\n".join(
            line for line in stderr.splitlines()
            if not line.strip().startswith("ERROR conda.cli.main_run")
        )
        return stdout, filtered_stderr
    except subprocess.TimeoutExpired as e:
        return "", "TimeoutError: 'Timed Out'"
    except Exception as e:
        return "", e
    finally:
        os.remove(tmp_file_path) 

class Interpreter:
    def __init__(
        self,
        timeout_length: int = 5,
    ) -> None:
        self.timeout_length = timeout_length
        self.codes = []

    def extra_code(self, code):
        match = re.search(r'```python\s*(.*?)```', code, re.DOTALL)
        if match:
            extra_code = match.group(1).strip()
        else:
            extra_code = code

        return extra_code

    def process_generation_to_code(self, gen):
        code = self.extra_code(gen)
        g_split = code.split("\n")
        run_code = self.codes.copy()
        for c in g_split:
            run_code.append(c)

        return run_code

    @staticmethod
    def execute(
        code,
        timeout_length=10,
    ):
        try:
            result, report = run_task(code, timeout_length)
            result = str(result).strip()
            report = str(report).strip()
            if not report:
                report = "Done"
                pickle.dumps(result)  # serialization check
        except:
            result = ""
            report = traceback.format_exc().split("\n")[-2]

        if 'No module named' in report:
            match = re.search(r"No module named ['\"]([^'\"]+)['\"]", report)
            if match:
                missing_full_module = match.group(1)
                top_module = missing_full_module.split('.')[0]

                try:
                    __import__(top_module)
                    print(f"[Interpreter warning] Module '{top_module}' is installed, but '{missing_full_module}' not found. Skip auto-install.")
                except ImportError:
                    if top_module.lower() == 'sklearn':
                        top_module = 'scikit-learn'
                    pip_res = subprocess.run(
                        ["/root/anaconda3/bin/conda", "run", "-n", "darl-run", "pip", "install", top_module], 
                        capture_output=True
                        )
                    if pip_res.returncode != 0:
                        print(f"[Interpreter pip error] an error occured when using pip to install missing package :{pip_res.stderr.decode('utf-8')}")
                    else:
                        print(f"[Interpreter pip] successfully pip install module: {top_module} restart the code")
                        try:
                            result, report = run_task(code, timeout_length)
                            result = str(result).strip()
                            report = str(report).strip()
                            if not report:
                                report = "Done"
                                pickle.dumps(result)  # serialization check
                        except:
                            result = ""
                            report = traceback.format_exc().split("\n")[-2]
        return result, report

    def apply(self, code):
        return self.run_code([code])[0]

    @staticmethod
    def truncate(s, max_length=400):
        half = max_length // 2
        if len(s) > max_length:
            s = s[:half] + "..." + s[-half:]
        return s

    def run_code(self, code):
        all_code_snippets = self.process_generation_to_code(code)

        try:
            result, report = self.execute(all_code_snippets, timeout_length=self.timeout_length)
        except TimeoutError as error:
            result = ""
            report = "Timeout Error"

        res, report = str(result).strip(), str(report).strip()
        res, report = self.truncate(res), self.truncate(report)
        self.update_code(report, code)

        return res, report
    
    def update_code(self, report, code):
        if report == 'Done':
            extra_code = self.extra_code(code)
            g_split = extra_code.split("\n")
            for c in g_split:
                if not (c.strip().startswith("print") or "print(" in c):
                    self.codes.append(c)
                else:
                    indent = re.match(r"^\s*",c).group(0)
                    self.codes.append(f"{indent}pass")
