import faulthandler
import itertools
import json
import multiprocessing
import os
import platform
import re
# to run the solution files we're using a timing based approach
import signal
import sys
from collections import defaultdict
# used for debugging to time steps
from datetime import datetime
from enum import Enum
# for capturing the stdout
from io import StringIO
# used for testing the code that reads from input
from unittest.mock import mock_open, patch

import numpy as np
from datasets import Dataset, DatasetDict, load_from_disk

try:
    from pyext import RuntimeModule
except ImportError:
    RuntimeModule = None

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

from .base import BaseDataset

TIMEOUT = 10


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

    @staticmethod
    def load(path: str, num_repeats: int = 1):
        dataset = load_from_disk(path)
        new_dataset = DatasetDict()
        # add new column "starter" in the prompt
        for split in dataset.keys():
            new_samples = []
            for idx, sample in enumerate(dataset[split]):
                starter_code = None if len(
                    sample['starter_code']) == 0 else sample['starter_code']
                try:
                    input_output = json.loads(sample['input_output'])
                    fn_name = (None if not input_output.get('fn_name') else
                               input_output['fn_name'])
                except ValueError:
                    fn_name = None
                starter = ''
                if starter_code:
                    starter += starter_code
                if (not fn_name) and (not starter_code):
                    call_format = '\\nUse Standard Input format'
                    starter += call_format
                else:
                    call_format = '\\nUse Call-Based format'
                    starter += call_format
                # Add the new column "starter" to the sample
                problem_id = f'{split}{idx}'
                sample['starter'] = starter
                sample['problem_id'] = problem_id
                new_samples.append(sample)
            new_data = {
                key: [sample[key] for sample in new_samples]
                for key in new_samples[0].keys()
            }
            new_dataset[split] = Dataset.from_dict(new_data)

        # num_repeats duplicate
        # train_repeated = []
        test_repeated = []
        # for sample in new_dataset['train']:
        #     train_repeated.extend([sample] * num_repeats)
        for sample in new_dataset['test']:
            test_repeated.extend([sample] * num_repeats)

        # dataset_train_repeated = new_dataset['train'].from_list(
        #     train_repeated
        # )
        dataset_test_repeated = new_dataset['test'].from_list(test_repeated)

        return DatasetDict({
            # 'train': dataset_train_repeated,
            'test': dataset_test_repeated
        })


EOF_STRINGS = ['\nQUESTION', '\n---', '\nANSWER', '<|endoftext|>']


@ICL_EVALUATORS.register_module()
class TACOEvaluator(BaseEvaluator):

    def post_process(self, text):
        if '```' in text:
            blocks = re.findall(r'```(.*?)```', text, re.DOTALL)
            if len(blocks) == 0:
                text = text.split('```')[1]  # fall back to default strategy
            else:
                text = blocks[0]  # fetch the first code block
                if not text.startswith('\n'):  # starting with ```python
                    text = text[max(text.find('\n') + 1, 0):]
        return text

    TIMEOUT = 10

    def check_correctness(self, sample, generation, timeout, debug=True):
        """Check correctness of code generation with a global timeout.

        The global timeout is to catch some extreme/rare cases not handled by
        the timeouts inside `run_test`
        """

        def _temp_run(sample, generation, debug, result):
            result.append(run_test(sample, test=generation, debug=debug))

        manager = multiprocessing.Manager()
        result = manager.list()
        p = multiprocessing.Process(target=_temp_run,
                                    args=(sample, generation, debug, result))
        p.start()
        p.join(timeout=timeout + 1)
        if p.is_alive():
            p.kill()
        if not result:
            in_outs = json.loads(sample['input_output'])
            # consider that all tests failed
            result = [[-1 for i in range(len(in_outs['inputs']))]]
            if debug:
                print('global timeout')
        return result[0]

    def evaluate_generations(self,
                             generations,
                             samples,
                             idx=None,
                             debug=False):
        assert len(generations.keys()) == len(samples)
        results = {}
        idx = 0
        for task_id, problem_generations in generations.items():
            sample = samples.iloc[idx]
            res = []
            # loop over the generations
            for o_idx, o in enumerate(problem_generations):
                curr_res = [-2]
                try:
                    curr_res = self.check_correctness(sample,
                                                      o,
                                                      timeout=TIMEOUT,
                                                      debug=debug)
                    if debug:
                        print(f'\nSuccessful compilation of task {o_idx}!')
                    fixed = []
                    for e in curr_res:
                        if isinstance(e, np.ndarray):
                            e = e.item(0)
                        if isinstance(e, np.bool_):
                            e = bool(e)
                        fixed.append(e)
                    curr_res = fixed
                    if not np.all(curr_res):
                        if debug:
                            print('Results were not True for all test cases')
                except Exception as e:
                    if debug:
                        print(
                            f'Compilation failed, test framework exception = {repr(e)}{e}\n'  # noqa: E501
                        )
                    break
                finally:
                    assert isinstance(curr_res, list)
                    res.append(curr_res)
            results[task_id] = res
            idx += 1
        return results

    def estimate_pass_at_k(self, num_samples, num_correct, k):
        """Estimates pass@k of each problem and returns them in an array."""

        def estimator(n: int, c: int, k: int) -> float:
            """Calculates 1 - comb(n - c, k) / comb(n, k)."""
            if n - c < k:
                return 1.0
            return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))

        if isinstance(num_samples, int):
            num_samples_it = itertools.repeat(num_samples, len(num_correct))
        else:
            assert len(num_samples) == len(num_correct)
            num_samples_it = iter(num_samples)

        return np.array([
            estimator(int(n), int(c), k)
            for n, c in zip(num_samples_it, num_correct)
        ])

    def compute_metrics(self, results, k_list=[1, 10, 100]):
        total = []
        correct = []
        task_ids = []
        for task_id, res in results.items():
            all_correct = []
            for generation in res:
                gen = np.array(generation)
                all_correct.append(np.all(gen > 0))
            task_ids.append(task_id)
            total.append(len(all_correct))
            correct.append(sum(all_correct))
        total = np.array(total)
        correct = np.array(correct)
        ks = k_list
        detail_pass_at_k = {
            f'pass@{k}': self.estimate_pass_at_k(total, correct, k).tolist()
            for k in ks if (total >= k).all()
        }
        pass_at_k = {
            f'pass@{k}':
            100 * (self.estimate_pass_at_k(total, correct, k).mean())
            for k in ks if (total >= k).all()
        }
        detail_metrics = {
            k: dict(zip(task_ids, v))
            for k, v in detail_pass_at_k.items()
        }
        pass_at_k['detail'] = detail_metrics
        return pass_at_k

    def score(self, predictions, references, test_set):
        assert len(predictions) == len(references)
        generations = defaultdict(list)
        for refer, pred in zip(references, predictions):
            pred = self.post_process(pred)
            generations[refer].append(pred)
        # convert to non-duplicated version
        test_set = test_set.to_pandas()
        test_set_origin = test_set.drop_duplicates(subset='problem_id')
        results = self.evaluate_generations(generations, test_set_origin)
        metrics = self.compute_metrics(results)
        return metrics


class CODE_TYPE(Enum):
    call_based = 0
    standard_input = 1


# stuff for setting up signal timer
class TimeoutException(Exception):
    pass


def timeout_handler(signum, frame):
    print('alarm went off')

    raise TimeoutException


try:
    signal.signal(signal.SIGALRM, timeout_handler)
except AttributeError:
    print('signal.SIGALRM is not available on this platform')
timeout = 4  # seconds


# used to capture stdout as a list
# from https://stackoverflow.com/a/16571630/6416660
# alternative use redirect_stdout() from contextlib
class Capturing(list):

    def __enter__(self):
        self._stdout = sys.stdout
        sys.stdout = self._stringio = StringIO()
        # Make closing the StringIO a no-op
        self._stringio.close = lambda x: 1
        return self

    def __exit__(self, *args):
        self.extend(self._stringio.getvalue().splitlines())
        del self._stringio  # free up some memory
        sys.stdout = self._stdout


def run_test(sample, test=None, debug=False):
    """if test(generated_code) is not None it'll try to run the code.

    otherwise it'll just return an input and output pair.
    """
    # Disable functionalities that can make destructive changes to the test.
    reliability_guard()

    if debug:
        print(f'start = {datetime.now().time()}')

    try:
        in_outs = json.loads(sample['input_output'])
    except ValueError:
        in_outs = None
    if in_outs:
        if in_outs.get('fn_name') is None:
            which_type = CODE_TYPE.standard_input  # Standard input
            method_name = None
        else:
            which_type = CODE_TYPE.call_based  # Call-based
            method_name = in_outs['fn_name']

    if debug:
        print(f'loaded input_output = {datetime.now().time()}')

    if test is None:
        return in_outs
    elif test is not None:
        results = []
        sol = 'import sys\nimport time\nimport itertools\nfrom itertools import accumulate, product, permutations, combinations\nimport collections\nfrom collections import Counter, OrderedDict, deque, defaultdict, ChainMap\nfrom functools import lru_cache\nimport math\nfrom math import sqrt, sin, cos, tan, ceil, fabs, floor, gcd, exp, log, log2\nimport fractions\nfrom typing import List, Tuple\nimport numpy as np\nimport random\nimport heapq\nfrom heapq import *\n'  # noqa: E501
        if debug:
            print(f'loading test code = {datetime.now().time()}')

        if which_type == CODE_TYPE.call_based:
            sol += test
            if debug:
                print(f'sol = {sol}')
            signal.alarm(timeout)
            try:
                tmp_sol = RuntimeModule.from_string('tmp_sol', '', sol)
                if 'class Solution' not in test:
                    tmp = tmp_sol
                else:
                    tmp = tmp_sol.Solution()
                signal.alarm(0)
            except Exception as e:
                signal.alarm(0)
                if debug:
                    print(f'type 0 compilation error = {e}')
                results.append(-2)
                return results
            signal.alarm(0)

        elif which_type == CODE_TYPE.standard_input:
            # sol
            tmp_test = test.split('\n')

            new_test = []
            for x in tmp_test:
                if (not x.startswith('from ')) and (
                        not x.startswith('import ')):
                    new_test.append('\t' + x + '\n')
                else:
                    new_test.append(x + '\n')
            tmp_test = new_test

            new_test = ''
            started = False
            for i in tmp_test:
                if i.startswith('\t') and not started:
                    new_test += 'stdin = sys.stdin\nstdout = sys.stdout\n'
                    new_test += 'def code():\n'
                    new_test += i
                    started = True
                elif started and ((i.startswith('from ')) or
                                  (i.startswith('import '))):
                    new_test += '\t' + i
                else:
                    new_test += i
            tmp_test = new_test

            sol += tmp_test
            if debug:
                print(f'sol = {sol}')
            method_name = 'code'
            signal.alarm(timeout)
            try:
                tmp_sol = RuntimeModule.from_string('tmp_sol', '', sol)
                tmp = tmp_sol
                signal.alarm(0)
            except Exception as e:
                signal.alarm(0)
                if debug:
                    print(f'type 1 compilation error = {e}')
                results.append(-2)
                return results
            signal.alarm(0)
        if debug:
            print(f'get method = {datetime.now().time()}')

        try:
            method = getattr(tmp,
                             method_name)  # get_attr second arg must be str
        except:  # noqa: E722
            signal.alarm(0)
            e = sys.exc_info()
            print(f'unable to get function error = {e}')
            results.append(-2)
            return results

        for index, inputs in enumerate(in_outs['inputs']):
            # JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list) # noqa: E501
            try:
                if isinstance(inputs[0], dict):
                    inputs = [{int(k): v for k, v in inputs[0].items()}]
            except:  # noqa: E722
                True
            try:
                if isinstance(in_outs['outputs'][index], dict):
                    in_outs['outputs'][index] = [{
                        int(k): v
                        for k, v in in_outs['outputs'][index].items()
                    }]
            except:  # noqa: E722
                True
            try:
                if isinstance(in_outs['outputs'][index][0], dict):
                    in_outs['outputs'][index] = [{
                        int(k): v
                        for k, v in in_outs['outputs'][index][0].items()
                    }]
            except:  # noqa: E722
                True

            if debug:
                print(
                    f'time: {datetime.now().time()} testing index = {index}  inputs = {inputs}, {type(inputs)}. type = {which_type}'  # noqa: E501
                )
            if which_type == CODE_TYPE.call_based:  # Call-based
                signal.alarm(timeout)
                faulthandler.enable()
                try:
                    output = method(*inputs)

                    # ground truth sequences are not tuples
                    if isinstance(output, tuple):
                        output = list(output)

                    tmp_result = output == in_outs['outputs'][index]
                    if isinstance(in_outs['outputs'][index],
                                  list) and in_outs['outputs'][index]:
                        tmp_result = tmp_result or (
                            output == in_outs['outputs'][index][0])

                    # ground truth sequences are not tuples
                    try:
                        if isinstance(output[0], tuple):
                            tmp_result = tmp_result or ([
                                list(x) for x in output
                            ] == in_outs['outputs'][index][0])
                    except:  # noqa: E722
                        True
                    results.append(tmp_result)

                    # reset the alarm
                    signal.alarm(0)
                except Exception as e:
                    signal.alarm(0)
                    faulthandler.disable()
                    if debug:
                        print(
                            f'Standard input runtime error or time limit exceeded error = {e}'  # noqa: E501
                        )
                    results.append(-1)
                    continue
                faulthandler.disable()
                signal.alarm(0)
                if debug:
                    print(
                        f"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}"  # noqa: E501
                    )
            elif which_type == CODE_TYPE.standard_input:  # Standard input
                faulthandler.enable()
                signal.alarm(timeout)
                passed = False

                if isinstance(inputs, list):
                    inputs = '\n'.join(inputs)
                if isinstance(in_outs['outputs'][index], list):
                    in_outs['outputs'][index] = '\n'.join(
                        in_outs['outputs'][index])

                with Capturing() as output:
                    try:
                        call_method(method, inputs)
                        # reset the alarm
                        signal.alarm(0)
                        passed = True
                    except Exception as e:
                        # runtime error or took too long
                        signal.alarm(0)
                        print(
                            f'Call-based runtime error or time limit exceeded error = {repr(e)}{e}'  # noqa: E501
                        )
                        results.append(-1)
                    signal.alarm(0)

                if not passed:
                    if debug:
                        nl = '\n'
                        if not isinstance(inputs, list):
                            print(
                                f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}"  # noqa: E501
                            )
                        else:
                            print(
                                f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}"  # noqa: E501
                            )
                    continue

                if passed and debug:
                    print(
                        f"==> output = {output}, test outputs = {in_outs['outputs'][index]}"  # noqa: E501
                    )

                if custom_compare_(output, in_outs['outputs'][index]):
                    tmp_result = True
                    results.append(tmp_result)
                    continue

                # ground truth sequences are expressed as lists not tuples
                if isinstance(output, tuple):
                    output = list(output)

                tmp_result = False
                try:
                    tmp_result = (output == [in_outs['outputs'][index]])
                    if isinstance(in_outs['outputs'][index], list):
                        tmp_result = tmp_result or (
                            output == in_outs['outputs'][index])
                        if isinstance(output[0], str):
                            tmp_result = tmp_result or ([
                                e.strip() for e in output
                            ] == in_outs['outputs'][index])
                except Exception as e:
                    if debug:
                        print(f'Failed check1 exception = {e}')
                    pass

                if tmp_result == True:  # noqa: E712
                    results.append(tmp_result)
                    continue

                # try one more time without \n
                if isinstance(in_outs['outputs'][index], list):
                    for tmp_index, i in enumerate(in_outs['outputs'][index]):
                        in_outs['outputs'][index][tmp_index] = i.split('\n')
                        in_outs['outputs'][index][tmp_index] = [
                            x.strip()
                            for x in in_outs['outputs'][index][tmp_index] if x
                        ]
                else:
                    in_outs['outputs'][index] = in_outs['outputs'][
                        index].split('\n')
                    in_outs['outputs'][index] = list(
                        filter(len, in_outs['outputs'][index]))
                    in_outs['outputs'][index] = list(
                        map(lambda x: x.strip(), in_outs['outputs'][index]))

                try:
                    tmp_result = (output == [in_outs['outputs'][index]])
                    if isinstance(in_outs['outputs'][index], list):
                        tmp_result = tmp_result or (
                            output == in_outs['outputs'][index])
                except Exception as e:
                    if debug:
                        print(f'Failed check2 exception = {e}')
                    pass

                if tmp_result == True:  # noqa: E712
                    results.append(tmp_result)
                    continue

                # try by converting the output into a split up list too
                if isinstance(output, list):
                    output = list(filter(len, output))

                if debug:
                    nl = '\n'
                    if not isinstance(inputs, list):
                        print(
                            f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}"  # noqa: E501
                        )
                    else:
                        print(
                            f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}"  # noqa: E501
                        )

                if tmp_result == True:  # noqa: E712
                    results.append(tmp_result)
                    continue

                try:
                    tmp_result = (output == [in_outs['outputs'][index]])
                    if isinstance(in_outs['outputs'][index], list):
                        tmp_result = tmp_result or (
                            output == in_outs['outputs'][index])
                except Exception as e:
                    if debug:
                        print(f'Failed check3 exception = {e}')
                    pass

                try:
                    output_float = [float(e) for e in output]
                    gt_float = [float(e) for e in in_outs['outputs'][index]]
                    tmp_result = tmp_result or (
                        (len(output_float) == len(gt_float))
                        and np.allclose(output_float, gt_float))
                except Exception:
                    pass
                try:
                    if isinstance(output[0], list):
                        output_float = [float(e) for e in output[0]]
                        gt_float = [
                            float(e) for e in in_outs['outputs'][index][0]
                        ]
                        tmp_result = tmp_result or (
                            (len(output_float) == len(gt_float))
                            and np.allclose(output_float, gt_float))
                except Exception:
                    pass

                if tmp_result == True:  # noqa: E712
                    results.append(tmp_result)
                    continue

                # try by converting the stuff into split up list
                if isinstance(in_outs['outputs'][index], list):
                    for tmp_index, i in enumerate(in_outs['outputs'][index]):
                        in_outs['outputs'][index][tmp_index] = set(i.split())
                else:
                    in_outs['outputs'][index] = set(
                        in_outs['outputs'][index].split())

                try:
                    tmp_result = (output == in_outs['outputs'][index])
                except Exception as e:
                    if debug:
                        print(f'Failed check4 exception = {e}')
                    continue

                if tmp_result == True:  # noqa: E712
                    results.append(tmp_result)
                    continue

                # try by converting the output into a split up list too
                if isinstance(output, list):
                    for tmp_index, i in enumerate(output):
                        output[tmp_index] = i.split()
                    output = list(filter(len, output))
                    for tmp_index, i in enumerate(output):
                        output[tmp_index] = set(i)
                else:
                    output = output.split()
                    output = list(filter(len, output))
                    output = set(output)

                try:
                    tmp_result = (set(frozenset(s) for s in output) == set(
                        frozenset(s) for s in in_outs['outputs'][index]))
                except Exception as e:
                    if debug:
                        print(f'Failed check5 exception = {e}')

                # if they are all numbers, round so that similar numbers are treated as identical # noqa: E501
                try:
                    tmp_result = tmp_result or (
                        set(
                            frozenset(round(float(t), 3) for t in s)
                            for s in output) ==  # noqa: E502, E128
                        set(
                            frozenset(round(float(t), 3) for t in s)
                            for s in in_outs['outputs'][index]))  # noqa: E501
                except Exception as e:
                    if debug:
                        print(f'Failed check6 exception = {e}')

                if tmp_result == True and debug:  # noqa: E712
                    print('PASSED')

                results.append(tmp_result)

                if debug:
                    nl = '\n'
                    if not isinstance(inputs, list):
                        print(
                            f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}"  # noqa: E501
                        )
                    else:
                        print(
                            f"output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}"  # noqa: E501
                        )

    return results


def custom_compare_(output, ground_truth):

    if isinstance(output, list):
        output_1 = '\n'.join(output)
        if stripped_string_compare(output_1, ground_truth):
            return True

    if isinstance(output, list):
        output_2 = [o.lstrip().rstrip() for o in output]
        output_2 = '\n'.join(output_2)
        if stripped_string_compare(output_2, ground_truth):
            return True

    return False


def stripped_string_compare(s1, s2):
    s1 = s1.lstrip().rstrip()
    s2 = s2.lstrip().rstrip()
    return s1 == s2


def call_method(method, inputs):

    if isinstance(inputs, list):
        inputs = '\n'.join(inputs)

    inputs_line_iterator = iter(inputs.split('\n'))

    # sys.setrecursionlimit(10000)

    # @patch('builtins.input', side_effect=inputs.split("\n"))
    @patch('builtins.open', mock_open(read_data=inputs))
    @patch('sys.stdin', StringIO(inputs))
    @patch('sys.stdin.readline', lambda *args: next(inputs_line_iterator))
    @patch('sys.stdin.readlines', lambda *args: inputs.split('\n'))
    @patch('sys.stdin.read', lambda *args: inputs)
    # @patch('sys.stdout.write', print)
    def _inner_call_method(_method):
        try:
            return _method()
        except SystemExit:
            pass
        finally:
            pass

    return _inner_call_method(method)


def reliability_guard(maximum_memory_bytes=None):
    """This disables various destructive functions and prevents the generated
    code from interfering with the test (e.g. fork bomb, killing other
    processes, removing filesystem files, etc.) WARNING This function is NOT a
    security sandbox.

    Untrusted code, including, model- generated code, should not be blindly
    executed outside of one. See the Codex paper for more information about
    OpenAI's code sandbox, and proceed with caution.
    """

    if maximum_memory_bytes is not None:
        import resource

        resource.setrlimit(resource.RLIMIT_AS,
                           (maximum_memory_bytes, maximum_memory_bytes))
        resource.setrlimit(resource.RLIMIT_DATA,
                           (maximum_memory_bytes, maximum_memory_bytes))
        if not platform.uname().system == 'Darwin':
            resource.setrlimit(resource.RLIMIT_STACK,
                               (maximum_memory_bytes, maximum_memory_bytes))

    faulthandler.disable()

    import builtins

    builtins.exit = None
    builtins.quit = None

    os.environ['OMP_NUM_THREADS'] = '1'

    os.kill = None
    os.system = None
    os.putenv = None
    os.remove = None
    os.removedirs = None
    os.rmdir = None
    os.fchdir = None
    os.setuid = None
    os.fork = None
    os.forkpty = None
    os.killpg = None
    os.rename = None
    os.renames = None
    os.truncate = None
    os.replace = None
    os.unlink = None
    os.fchmod = None
    os.fchown = None
    os.chmod = None
    os.chown = None
    os.chroot = None
    os.fchdir = None
    os.lchflags = None
    os.lchmod = None
    os.lchown = None
    os.getcwd = None
    os.chdir = None

    import shutil

    shutil.rmtree = None
    shutil.move = None
    shutil.chown = None

    import subprocess

    subprocess.Popen = None  # type: ignore

    __builtins__['help'] = None

    import sys

    sys.modules['ipdb'] = None
    sys.modules['joblib'] = None
    sys.modules['resource'] = None
    sys.modules['psutil'] = None
    sys.modules['tkinter'] = None
