# flake8: noqa
import json
import os.path as osp
import re

from datasets import Dataset

from opencompass.openicl.icl_evaluator import BaseEvaluator
from opencompass.registry import LOAD_DATASET
from opencompass.utils import get_data_path

from ..base import BaseDataset


def check_match(target, generation):
    pattern_str = target.replace('{{',
                                 '{').replace('}}',
                                              '}').replace('{answer}', '.*')
    pattern_str = re.escape(pattern_str).replace('\\.\\*', '.*')
    match = re.fullmatch(pattern_str, generation)
    return bool(match)


@LOAD_DATASET.register_module()
class FollowBenchDataset(BaseDataset):

    def load(self, path: str, name: str, cate: str, *args, **kwargs):

        path = get_data_path(path, local_mode=True)
        filename = osp.join(path, f'{name}.json')
        raw_data = []
        raw_data = []
        with open(filename, 'r', encoding='utf-8') as f:
            data = json.load(f)
            for item in data:
                if cate == 'llm':
                    raw_data.append({
                        'instruction': item['instruction'],
                        'judge_prompt': item['judge_prompt'],
                        'judge': item
                    })
                elif cate == 'rule':
                    raw_data.append({
                        'instruction': item['instruction'],
                        'judge': item
                    })
                else:
                    raise NotImplementedError(
                        f"Category '{cate}' is not implemented.")

        dataset = Dataset.from_list(raw_data)
        return dataset


class FollowBenchEvaluator(BaseEvaluator):
    """Evaluator for followbench rule-based 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):
        results = {'example': {'accuracy': [0, 0, 0, 0, 0], 'num': 0}}
        for prediction, reference in zip(predictions, references):
            if reference['category'] == 'example':
                results['example']['num'] += 1
                template = reference['target'].replace('{instruction}\n', '')
                match_result = check_match(template, prediction)
                print(match_result)
                results['example']['accuracy'][reference['level'] -
                                               1] += match_result
        results['example']['accuracy'] = [
            round(acc / (results['example']['num'] // 5), 2)
            for acc in results['example']['accuracy']
        ]
        ######## Still not finished for rule-based evaluation

        # 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}