# borrowed and extended from
# https://github.com/Naman-ntc/codescratch/blob/main/evaluation/bigcode-evaluation-harness/lm_eval/tasks/custom_metrics/apps_custom_metrics/utils.py

import os

from livebench.lcb_runner.evaluation.base import LCBCorrectnessCriteria
from livebench.lcb_runner.evaluation.java import JavaCorrectnessEvaluator

os.environ["TOKENIZERS_PARALLELISM"] = "false"
import json
import multiprocessing
from concurrent.futures import ProcessPoolExecutor, as_completed


import numpy as np
from tqdm import tqdm
from pathlib import Path

from livebench.lcb_runner.evaluation.testing_util import run_test
from livebench.lcb_runner.evaluation.pass_k_utils import compute_metrics_from_results
from livebench.common import ProgrammingLanguage

def _temp_run(sample, generation, language, debug, result, metadata_list, timeout):
    res, metadata = run_test(sample, language, test=generation, debug=debug, timeout=timeout)
    result.append(res)
    metadata_list.append(metadata)


def check_correctness(sample, generation, language, 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`"""

    manager = multiprocessing.Manager()
    result = manager.list()
    metadata_list = manager.list()
    p = multiprocessing.Process(
        target=_temp_run,
        args=(sample, generation, language, debug, result, metadata_list, timeout),
    )
    p.start()
    p.join(
        timeout=(timeout + 1) * len(json.loads(sample["input_output"])["inputs"]) + 5
    )
    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(f"global timeout")
    return result[0], metadata_list[0]


def evaluate_generations_by_problem(args):
    problem_generations: list[str] = args[0]
    sample = args[1]
    language = args[2]
    debug: bool = args[3]
    timeout: int = args[4]
    res = []
    metadata = []
    for o_idx, o in enumerate(problem_generations):
        curr_res = [-2]
        try:
            if language == ProgrammingLanguage.JAVA:
                evaluator = JavaCorrectnessEvaluator(Path('data/live_bench/coding_java/LCB_generation/work'))
                curr_res, curr_metadata = evaluator.check(
                    sample,
                    o,
                    timeout=timeout,
                    debug=debug,
                )
            else:
                curr_res, curr_metadata = check_correctness(
                    sample, o, language, 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(f"Results were not True for all test cases {curr_res=}\n")
        except Exception as e:
            print(f"Compilation failed, test framework exception = {repr(e)}{e}\n")
            # break
            curr_metadata = {}
        finally:
            assert isinstance(curr_res, list)
            assert isinstance(curr_metadata, dict)
            res.append(curr_res)
            metadata.append(curr_metadata)
    if debug:
        for i, r in enumerate(problem_generations):
            print("Sample\n")
            print(r)
            print("\n")
            print("Result\n")
            print(res[i])
            print("*" * 30 + "\n\n")
    
    return res, metadata


def evaluate_generations(
    samples_list: list,
    generations_list: list[list[str]],
    language,
    debug: bool = False,
    num_process_evaluate: int = 16,
    timeout=6,
):
    """We take the list of code generations and try to compile them
     and the run their corresponding unit tests which are retrieved from the APPS dataset.

    Args:
        generations: list of code generations (same order as samples in APPS dataset)
        level: difficulty level used in the generation, can be "all", "introductory", "interview" or "competition"

    Returns:
        results: dictionary of results, key is the problem index, value is a list of results for each generation
        [-2] = compile error, [-1] = runtime error [False] = failed test case [True] = passed test case
    """
    # generations are code generations in the same order of the dataset

    inputs = [
        [(generations_list[index], samples_list[index], language, debug, timeout), index]
        for index in range(len(generations_list))
    ]
    if debug:
        with tqdm(total=len(inputs)) as pbar:
            with ProcessPoolExecutor(
                max_workers=1 if debug else num_process_evaluate
            ) as executor:
                futures = {
                    executor.submit(evaluate_generations_by_problem, arg): index
                    for arg, index in inputs
                }

                results = {}
                metadata = {}
                for future in as_completed(futures):
                    index = futures[future]
                    results[index], metadata[index] = future.result()
                    pbar.update(1)
    else:
        with ProcessPoolExecutor(
            max_workers=1 if debug else num_process_evaluate
        ) as executor:
            futures = {
                executor.submit(evaluate_generations_by_problem, arg): index
                for arg, index in inputs
            }

            results = {}
            metadata = {}
            for future in as_completed(futures):
                index = futures[future]
                results[index], metadata[index] = future.result()

    assert len(results) == len(
        inputs
    ), f"results = {len(results)} inputs = {len(inputs)} {results=}"
    # results = {i: r for r, (_, i) in zip(results, inputs)}

    return results, metadata


def codegen_metrics(
    samples,
    generations,
    language,
    k_list=[1, 5],
    num_process_evaluate=16,
    timeout=6,
    debug=False,
):
    results, metadata = evaluate_generations(
        samples,
        generations,
        language,
        debug=debug,
        num_process_evaluate=num_process_evaluate,
        timeout=timeout,
    )
    metrics = compute_metrics_from_results(results, k_list=k_list)

    final_metadata = []
    for key in sorted(list(metadata.keys())):
        final_metadata.append(metadata[key])
    for i in range(len(final_metadata)):
        if type(final_metadata[i]) is not list:
            final_metadata[i] = [json.dumps(final_metadata[i])]
        else:
            final_metadata[i] = [json.dumps(x) for x in final_metadata[i]]

        assert len(final_metadata[i]) == len(
            generations[0]
        ), f"{len(final_metadata[i])=}"

    return metrics, results, final_metadata
