"""
ClassEval: Class level Python code evaluation

Based on the paper https://arxiv.org/pdf/2308.01861 .
The datasets can be found either on https://huggingface.co/datasets/FudanSELab/ClassEval
or on https://github.com/FudanSELab/ClassEval

This is an inspect_ai implementation of the paper.
"""

import re
from typing import Any

from inspect_ai import Epochs, Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.scorer import (
    CORRECT,
    INCORRECT,
    Score,
    Scorer,
    Target,
    mean,
    scorer,
    std,
)
from inspect_ai.solver import (
    TaskState,
    generate,
    system_message,
)
from inspect_ai.util import ExecResult, sandbox

from .utils import construct_prompt

# Timeout for scoring.
VERIFY_TIMEOUT = 30


@task
def class_eval(few_shot: int = 1, few_shot_seed: int = 42) -> Task:
    """Inspect Task implementation of ClassEval.

    Args:
        few_shot (int): The number of few shots to include.
        few_shot_seed (int): The seed for generating few shots.
    """
    dataset = hf_dataset(
        path="FudanSELab/ClassEval",
        split="test",
        sample_fields=record_to_sample,
    )

    INSTRUCTION = """

    You are an expert Python programmer. You will be given a task, and the tests that your code must pass.

    """

    return Task(
        dataset=dataset,
        solver=[system_message(INSTRUCTION), generate()],
        scorer=class_eval_scorer(),
        epochs=Epochs(few_shot, [f"pass_at_{few_shot}"]),
        sandbox="docker",
    )


@scorer(metrics=[mean(), std()])
def class_eval_scorer() -> Scorer:
    """
    Scorer for class eval

    It will first identify the python code in the output of the model.
    Then we append the test cases to the code and execute it.
    If the code passes all the test cases, we return a CORRECT score.
    Otherwise, we return an INCORRECT score.
    """

    async def score(state: TaskState, target: Target) -> Score:
        generated_code = find_code(state.output.completion)
        code = generated_code + "\n" + state.metadata["test"]

        explanation = ""
        explanation += "The following code was executed:\n\n```python\n"
        explanation += code
        explanation += "\n```\n"

        try:
            result: ExecResult[str] = await sandbox().exec(
                cmd=["python", "-c", code],
                timeout=VERIFY_TIMEOUT,
            )

            if result.success:
                explanation += "All test cases passed.\n"
            else:
                explanation += "Code did not pass all test cases.\n"
                if result.stderr:
                    explanation += (
                        f"See details below.\n ```python\n {result.stderr} \n ```\n"
                    )
        except TimeoutError:
            result = ExecResult(False, 1, "", "Verification timed out.")
            explanation += "Verification timed out."

        return Score(
            value=CORRECT if result.success else INCORRECT,
            answer=generated_code,
            explanation=explanation,
        )

    return score


def find_code(completion: str) -> str:
    """Remove Markdown formatting around generated code blocks."""
    pattern = re.compile(r"```python\n(.*?)```", re.DOTALL)
    matches = pattern.findall(completion)
    extracted_answer = matches[0] if len(matches) >= 1 else completion

    return str(extracted_answer)


def record_to_sample(
    record: dict[str, Any],
) -> Sample:
    """Convert a record from the dataset to a Sample"""
    return Sample(
        input=construct_prompt(record),
        target=record["solution_code"],
        id=record["task_id"],
        metadata={
            "task_id": record["task_id"],
            "skeleton": record["skeleton"],
            "test": record["test"],
            "solution_code": record["solution_code"],
            "import_statement": record["import_statement"],
            "class_description": record["class_description"],
            "methods_info": record["methods_info"],
            "class_name": record["class_name"],
            "test_classes": record["test_classes"],
            "class_constructor": record["class_constructor"],
            "fields": record["fields"],
        },
    )
