import configparser
import importlib
import os
import pickle
import re
import shutil
import signal
import sys
import tempfile
import threading
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from subprocess import PIPE, Popen
from typing import Optional, Union

from datasets import Dataset, DatasetDict

from opencompass.openicl.icl_evaluator import BaseEvaluator
from opencompass.registry import LOAD_DATASET, TEXT_POSTPROCESSORS

from .base import BaseDataset


@LOAD_DATASET.register_module()
class DS1000Dataset(BaseDataset):
    """DS1000 dataset.

    The DS-1000 dataset can be downloaded from
    https://github.com/HKUNLP/DS-1000/blob/main/ds1000_data.zip
    """

    ALL_LIBS = [
        'Pandas',
        'Numpy',
        'Matplotlib',
        'Tensorflow',
        'Scipy',
        'Sklearn',
        'Pytorch',
    ]

    def get_data(self, problem_path: str) -> dict:
        """Get necessary data from problem path."""

        problem_path = Path(problem_path)
        data = dict()

        # load meta information in .cfg
        problem_config = configparser.RawConfigParser()
        problem_config.read(problem_path / '.cfg')
        for args in [
            ('library', 'lib'),
            ('test', 'test_type'),
            ('test', 'test_case_cnt'),
            ('perturbation', 'perturbation_type'),
            ('perturbation', 'perturbation_origin_id'),
        ]:
            data[args[1]] = problem_config.get(*args)

        # read problem content files
        for file_name in [
                'reference_code.txt',
                'test_code.py',
                'code_context.txt',
                'test_generate_pickle.py',
                'prompt.txt',
        ]:
            with open(problem_path / file_name, 'r', encoding='UTF-8') as f:
                data[file_name.split('.')[0]] = f.read()

        data['test_column'] = dict(
            problem_path=str(problem_path),
            code_context=data.pop('code_context'),
            test_case_cnt=data.pop('test_case_cnt'),
            test_type=data.pop('test_type'),
        )
        return data

    def load(self,
             path: str,
             libs: Optional[Union[str, list]] = None,
             mode: str = 'Insertion'):
        """Load dataset.

        Args:
            path(str): Data path of the dataset.
            libs(Optional[Union[str, list]]): All the libraries to be
                tested. Defaults to None, which means all.
            mode(str): Dataset mode. Only `Insertion` and `Completion`
                are supported. Defaults to `Insertion`.
        """

        if isinstance(libs, str):
            libs = [libs]

        if libs:
            for lib in libs:
                assert lib in self.ALL_LIBS, \
                    f'Please choose libs in {self.ALL_LIBS}'
        else:
            libs = self.ALL_LIBS

        data_list = []
        for lib in libs:
            source_path = Path(path) / lib / mode
            problems = sorted(os.listdir(source_path),
                              key=lambda x: int(str(x).replace('q', '')))
            for problem in problems:
                new_problem = self.get_data(source_path / problem)
                data_list.append(new_problem)

        dataset = Dataset.from_list(data_list)
        return DatasetDict({'test': dataset})


@TEXT_POSTPROCESSORS.register_module('ds1000')
def ds1000_postprocess(text: str) -> str:
    match = re.search('<code>(.*?)</code>', text, re.DOTALL)
    if match:
        text = match.group(1)

    match = re.search('```python(.*?)```', text, re.DOTALL)
    if match:
        text = match.group(1)

    match = re.search('```(.*?)```', text, re.DOTALL)
    if match:
        text = match.group(1)

    match = re.search('BEGIN SOLUTION(.*?)END SOLUTION', text, re.DOTALL)
    if match:
        text = match.group(1)

    return text


@TEXT_POSTPROCESSORS.register_module('ds1000_matplotlib')
def ds1000_matplotlib_postprocess(text: str) -> str:
    text = ds1000_postprocess(text)
    code_lines = text.split('\n')
    postprocessed_lines = []
    for line in code_lines:
        skip_line_flag = False
        # Matplotlib removes function calls that will clear the
        # global figure object
        # Removing these functions helps running execution-based evaluation
        for phrase in ['plt.show()', 'plt.clf()', 'plt.close()', 'savefig']:
            if phrase in line:
                skip_line_flag = True
                break

        if skip_line_flag:
            continue
        else:
            postprocessed_lines.append(line)
    text = '\n'.join(postprocessed_lines)
    return text


class DS1000Evaluator(BaseEvaluator):
    """Evaluator for ds1000 eval."""

    def __init__(self, num_workers=16) -> None:
        self.num_workers = num_workers

    def score_single(self, pred, refer):
        generated_code = pred

        # get current dir because we will enter a temp dir to
        # execute generated code
        cwd = os.getcwd()

        def chdir_return(cwd, return_value):
            os.chdir(cwd)
            return return_value

        # we create a tempdir to execute each generated program
        with tempfile.TemporaryDirectory() as tempdir_name:

            tempdir_name = Path(tempdir_name)
            # copy all files and data dependencies from
            shutil.copytree(refer['problem_path'],
                            tempdir_name,
                            dirs_exist_ok=True)
            # generated outputs will be put into `result`
            os.mkdir(tempdir_name / 'result')

            program = refer['code_context'].replace('[insert]', generated_code)
            with open(tempdir_name / 'program.py', 'w', encoding='UTF-8') as f:
                f.write(program)

            # enter into the tempdir to execute
            os.chdir(tempdir_name)

            execution_status = []
            # a question may not have test case but we can still execute and
            # see if there is error
            test_cnt = max(1, int(refer['test_case_cnt']))
            for i in range(1, test_cnt + 1):
                # notice this command, e.g., you may need to
                # replace `python` with `python3`
                cmd_text = f'python program.py --test_case {i}'
                time_limit = 60  # should not change the official time_limit
                cmd = Command(cmd_text, )
                exit_code = cmd.run(
                    timeout=time_limit)  # 0 if there is no error
                execution_status.append(exit_code)

            # return if execution error
            if sum(execution_status) > 0:
                return chdir_return(cwd, False)

            # loading testing code as a module
            test_module = import_source_file(tempdir_name / 'test_code.py',
                                             'test_code')
            pass_flag = True

            if int(refer['test_type']) == 3:
                # stringTest parses the generated code into AST
                # and check AST components
                # if there is static error, stringTest may raise an exception
                generated_code = generated_code.split('\n')
                for line in generated_code:
                    if 'print' in line and '#' not in line.split('print'):
                        generated_code.remove(line)
                generated_code = '\n'.join(generated_code)
                try:
                    pass_flag = test_module.stringTest(generated_code)
                except Exception:
                    # return False if stringTest error
                    return chdir_return(cwd, False)

            test_cnt = max(int(refer['test_case_cnt']), 1)
            for i in range(1, test_cnt + 1):
                try:
                    ans = pickle.load(open(f'ans/ans{i}.pkl', 'rb'))
                    # loading the generated output might still raise Exception
                    # if the generated code is not correct
                    result = pickle.load(open(f'result/result_{i}.pkl', 'rb'))
                    pass_flag = test_module.test(result, ans) == 1
                except Exception:
                    # return False if stringTest error
                    return chdir_return(cwd, False)

        return chdir_return(cwd, pass_flag)

    def score(self, predictions, references):
        total = len(predictions)

        # Each process changes cwd, need to use multi-processing
        with ProcessPoolExecutor(self.num_workers) as executor:
            passed = sum(
                list(executor.map(self.score_single, predictions, references)))

        return {'accuracy': passed / total}


class Command(object):
    """This object takes in command and executes it with time out."""

    def __init__(self, cmd):
        self.cmd = cmd
        self.process = None

    def run(self, timeout):

        def target():
            # Check if in Windows https://stackoverflow.com/questions/1325581/how-do-i-check-if-im-running-on-windows-in-python  # noqa
            if os.name == 'nt':
                self.process = Popen(self.cmd,
                                     shell=True,
                                     stdout=PIPE,
                                     stderr=PIPE)
            else:
                self.process = Popen(self.cmd,
                                     shell=True,
                                     stdout=PIPE,
                                     stderr=PIPE,
                                     preexec_fn=os.setsid)
            self.process.communicate()

        thread = threading.Thread(target=target)
        thread.start()

        thread.join(timeout)
        if thread.is_alive():
            # Check if in Windows https://stackoverflow.com/questions/1325581/how-do-i-check-if-im-running-on-windows-in-python  # noqa
            if os.name == 'nt':
                Popen('TASKKILL /F /PID {pid} /T'.format(pid=self.process.pid))
            else:
                os.killpg(self.process.pid, signal.SIGTERM)
            thread.join()
        return self.process.returncode


def import_source_file(fname, modname):
    """Import a Python source file and return the loaded module.

    Args:
        fname: The full path to the source file.  It may container characters
            like `.` or `-`.
        modname: The name for the loaded module.  It may contain `.` and even
            characters that would normally not be allowed (e.g., `-`).
    Return:
        The imported module

    Raises:
        ImportError: If the file cannot be imported (e.g, if it's not a `.py`
            file or if it does not exist).
        Exception: Any exception that is raised while executing the module
            (e.g. :exc:`SyntaxError).
            These are errors made by the author of the module!
    """
    # https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
    spec = importlib.util.spec_from_file_location(modname, fname)
    if spec is None:
        raise ImportError(
            f"Could not load spec for module '{modname}' at: {fname}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[modname] = module
    try:
        spec.loader.exec_module(module)
    except FileNotFoundError as e:
        raise ImportError(f'{e.strerror}: {fname}') from e
    return module