"""
BigCodeBench: Benchmarking Code Generation with Diverse Function Calls and Complex Instructions

Terry Yue Zhuo, Minh Chien Vu, Jenny Chim, Han Hu, Wenhao Yu, Ratnadira Widyasari, Imam Nur Bani Yusuf, Haolan Zhan, Junda He, Indraneil Paul, Simon Brunner, Chen Gong, Thong Hoang, Armel Randy Zebaze, Xiaoheng Hong, Wen-Ding Li, Jean Kaddour, Ming Xu, Zhihan Zhang, Prateek Yadav, Naman Jain, Alex Gu, Zhoujun Cheng, Jiawei Liu, Qian Liu, Zijian Wang, David Lo, Binyuan Hui, Niklas Muennighoff, Daniel Fried, Xiaoning Du, Harm de Vries, Leandro Von Werra

ArXiv: https://arxiv.org/abs/2406.15877
Website: https://bigcode-bench.github.io/
HF: https://huggingface.co/datasets/bigcode/bigcodebench
Github: https://github.com/bigcode-project/bigcodebench

Code inspired by [humaneval](../humaneval/humaneval.py).
"""

import re
from pathlib import Path
from typing import Any, Callable, List, Optional, Union
from warnings import warn

from datasets import load_dataset  # type: ignore
from inspect_ai import Task, task
from inspect_ai.dataset import Sample
from inspect_ai.scorer import (
    CORRECT,
    INCORRECT,
    Score,
    Scorer,
    Target,
    mean,
    scorer,
    std,
)
from inspect_ai.solver import Solver, TaskState, generate
from inspect_ai.util import ExecResult, sandbox

VERIFY_TIMEOUT = 30

INSTRUCTION_PROMPT = """Read the following function signature and docstring, and fully implement the function described. Make sure to include the import statements and use the same variable names as stated in the header.
"""


@task
def bigcodebench(
    solver: Optional[Solver] = None,
    instruction_prompt: str = INSTRUCTION_PROMPT,
    scorer: Union[Scorer, List[Scorer], None] = None,
    sandbox: str = "docker",
    version: str = "v0.1.2",
    subset: Optional[List[int]] = None,  # not needed for command line inspect eval
) -> Task:
    """Inspect implementation for the BigCodeBench benchmark

    Defaults to pulling v0.1.2 from huggingface
    Completion version of the prompts only
    See https://arxiv.org/abs/2406.15877 for more information about the dataset

    To evaluate on a subset of the dataset using the command line tool, use the --sample-id argument
    inspect eval inspect_evals/bigcodebench --sample-id=123,456

    However, the subset argument can be helpful if you want a task object with only a subset
    of the question.

    Args:
        solver (Solver): The solver to use for this evaluation.
        instruction_prompt (String): The prompt to prepend to the code problem.
        scorer (Scorer): The scorer to use for this evaluation.
        sandbox (String): The sandbox to use for this evaluation. Defaults to "docker".
        version (String): Version number for the evaluation from huggingface (https://huggingface.co/datasets/bigcode/bigcodebench)
        subset: (Optional[List[int]]): Evaluates on a subset of BigCodeBench, as opposed to the full set.
    """
    bigcode = load_dataset("bigcode/bigcodebench", split=version)
    record_to_sample = get_record_to_sample(instruction_prompt)
    if subset:
        dataset = [record_to_sample(bigcode[i]) for i in subset]
    else:
        dataset = [record_to_sample(record) for record in bigcode]
    return Task(
        dataset=dataset,
        solver=solver or generate(),
        scorer=scorer or verify(),
        sandbox=(sandbox, (Path(__file__).parent / "compose.yaml").as_posix()),
    )


def get_record_to_sample(
    instruction_prompt: str = INSTRUCTION_PROMPT,
) -> Callable[[dict[str, Any]], Sample]:
    """Converts dataset to inspect Samples

    This implements the completion version as opposed to the instruct version
    The completion version would have examples of how a function works
    that is absent in the instruct version.
    """

    def record_to_sample(record: dict[str, Any]) -> Sample:
        # Get the integer id from the end of the task_id, which is formatted as "BigCodeBench/0"
        task_id = int(record["task_id"].split("/")[-1])
        return Sample(
            id=task_id,
            input=instruction_prompt + record["complete_prompt"],
            target=record["complete_prompt"] + record["canonical_solution"],
            metadata={
                "test": record["test"],
                "libs": record["libs"],
                "task_id_metadata": record["task_id"],
            },
        )

    return record_to_sample


def find_code(completion: str) -> str:
    """
    Removes Markdown formatting around generated code blocks.

    Copied from humaneval
    """
    pattern_1 = re.compile(r"```python\n(.*?)```", re.DOTALL)
    pattern_2 = re.compile(r"```\n(.*?)```", re.DOTALL)
    matches = pattern_1.findall(completion) + pattern_2.findall(completion)
    if matches:
        extracted_answer = matches[0]
    else:
        extracted_answer = completion
    return str(extracted_answer)


@scorer(metrics=[mean(), std()])
def verify() -> Scorer:
    """Runs the provided test cases in a docker container.

    Just looks for a code block, instead of using bigcodebench's
    [code sanitization](https://github.com/bigcode-project/bigcodebench/blob/main/bigcodebench/sanitize.py)
    """

    async def score(state: TaskState, target: Target) -> Score:
        # extract answer from completion and format for verification
        answer = find_code(state.output.completion)
        code = [
            answer,
            "\n",
            state.metadata["test"],
            "\n",
            """
if __name__ == '__main__':
    unittest.main()""",
        ]

        # verify (returns error status if incorrect)
        try:
            result = await sandbox().exec(
                cmd=["python", "-c", "".join(code)],
                timeout=VERIFY_TIMEOUT,
            )
        except TimeoutError:
            result = ExecResult(False, 1, "", "Verification timed out.")

        if not result.success:
            if "ModuleNotFoundError" in result.stderr:
                warn(
                    f"There is a missing module during the docker execution: for task {state.metadata['task_id_metadata']}. The required libraries for this question are {state.metadata['libs']} \n This resulted in {result.stderr}"
                )

        explanation = "".join(
            ["The following verification code was executed:\n\n"]
            + ["```python\n\n"]
            + code
            + ["\n```\n"]
        )
        if result.success:
            explanation = explanation + "The submission was correct."
        else:
            explanation = (
                explanation
                + f"The submission was incorrect with the following error: \n\n{result.stderr}"
            )
        # return score w/ model answer + scoring details
        return Score(
            value=CORRECT if result.success else INCORRECT,
            answer=answer,
            explanation=explanation,
        )

    return score
