"""[∞Bench: Extending Long Context Evaluation Beyond 100K Tokens](https://aclanthology.org/2024.acl-long.814) (Zhang et al., ACL 2024)"""

import re
from typing import Any, Callable, Literal

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.model import get_model
from inspect_ai.scorer import Scorer, choice, exact, includes, pattern
from inspect_ai.solver import (
    Solver,
    generate,
    multiple_choice,
    prompt_template,
    system_message,
)

from inspect_evals.infinite_bench.constants import (
    CODE_RUN_ANSWER_PATTERN,
    MATH_FIND_INPUT_TO_PROMPT,
    TASK_METADATA,
    ft,
)
from inspect_evals.infinite_bench.math_calc_scorer import math_calc_scorer
from inspect_evals.infinite_bench.truncate_input import truncate
from inspect_evals.infinite_bench.utils import create_system_msg

# Retrieval tasks #


@task
def infinite_bench_kv_retrieval() -> Task:
    return infinite_bench_task(
        task_name="kv_retrieval",
        solver=generate(),
        scorer=includes(),
    )


@task
def infinite_bench_number_string() -> Task:
    return infinite_bench_task(
        task_name="number_string",
        solver=generate(),
        scorer=includes(),
    )


@task
def infinite_bench_passkey() -> Task:
    return infinite_bench_task(
        task_name="passkey",
        solver=generate(),
        scorer=includes(),
    )


# Coding tasks #


@task
def infinite_bench_code_debug() -> Task:
    return infinite_bench_task(
        task_name="code_debug",
        solver=multiple_choice(),
        scorer=choice(),
    )


@task
def infinite_bench_code_run() -> Task:
    return infinite_bench_task(
        task_name="code_run",
        solver=generate(),
        scorer=pattern(pattern=CODE_RUN_ANSWER_PATTERN),
    )


# Math tasks #


@task
def infinite_bench_math_calc() -> Task:
    model_name = get_model().name
    return infinite_bench_task(
        task_name="math_calc",
        solver=generate(),
        scorer=math_calc_scorer(model_name),
    )


@task
def infinite_bench_math_find() -> Task:
    return infinite_bench_task(
        task_name="math_find",
        solver=generate(),
        scorer=exact(),
    )


# Dialogue-based tasks #


# Note: the task prompt asks the model to identify a single character
# from a given dialogue. However, the dataset includes some samples
# where the target is a list of characters and the paper says "In the
# En.Dia task, random instances of character names within a script
# are replaced with $$MASK$$. The objective is to correctly identify
# these masked names." Therefore, we use includes() instead of
# exact().
@task
def infinite_bench_longdialogue_qa_eng() -> Task:
    return infinite_bench_task(
        task_name="longdialogue_qa_eng",
        solver=generate(),
        scorer=includes(),
    )


# Novel-based tasks #


@task
def infinite_bench_longbook_choice_eng() -> Task:
    return infinite_bench_task(
        task_name="longbook_choice_eng",
        solver=multiple_choice(),
        scorer=choice(),
    )


def infinite_bench_task(
    task_name: Literal[
        "passkey",
        "kv_retrieval",
        "number_string",
        "code_run",
        "code_debug",
        "math_find",
        "math_calc",
        "longdialogue_qa_eng",
        "longbook_choice_eng",
    ],
    solver: Solver,
    scorer: Scorer,
) -> Task:
    """Inspect task implementing the InfiniteBench benchmark.

    Args:
        task_name (Literal): The specific task to load, such as "passkey", "kv_retrieval", etc.
        solver (Solver): The solver to use for this task.
        scorer (Scorer): The scorer to use for evaluating this task.

    Returns:
        Task: The task with the corresponding solver and scorer.
    """
    # Get the task-specific metadata
    task_metadata = TASK_METADATA[task_name]

    # Use the factory function to get the appropriate sample function
    sample_function = get_sample_function(task_name)

    # Load the dataset for the specific task
    dataset = hf_dataset(
        path="xinrongzhang2022/InfiniteBench",
        name="default",
        data_files={
            task_name: task_metadata["data_file"]
        },  # Load the specific data file
        split=task_name,
        features=ft,
        sample_fields=sample_function,
    )

    # Create the system message for the task
    task_system_message = create_system_msg(task_name)

    # Retrieve the task-specific prompt template
    task_prompt_template = task_metadata["prompt_template"]

    # Return the task with the corresponding solver and scorer
    return Task(
        dataset=dataset,
        solver=[
            system_message(task_system_message),  # Inject the system message
            prompt_template(
                task_prompt_template
            ),  # Add the task-specific prompt template here
            solver,  # Pass the solver provided
        ],
        scorer=scorer,  # Pass the scorer provided
    )


def get_sample_function(task_name: str) -> Callable[[dict[str, Any]], Sample]:
    """
    Return the appropriate record_to_sample function based on the task name.

    Args:
        task_name (str): The name of the task to get the sample function for.

    Returns:
        Callable: The record_to_sample function for the task.

    Raises:
        ValueError: If the task_name does not match any known tasks.
    """
    if task_name in ["kv_retrieval", "number_string", "passkey"]:
        return record_to_sample_single_element_answer
    elif task_name in ["longdialogue_qa_eng", "math_calc"]:
        return record_to_sample_multi_element_answer
    elif task_name in ["code_debug", "longbook_choice_eng"]:
        return record_to_sample_multiple_choice
    elif task_name == "code_run":
        return record_to_sample_code_run
    elif task_name == "math_find":
        return record_to_sample_math_find
    else:
        raise ValueError(f"Unknown task name: {task_name}")


def record_to_sample_single_element_answer(record: dict[str, Any]) -> Sample:
    """
    Convert a record from the dataset to a sample for tasks with a single element answer.

    Args:
        record (dict): The record from the dataset to convert to a sample.

    Returns:
        Sample: The sample with the corresponding id, input, and target.
    """
    # Construct the input string from context and input
    input_str = f"{record['context']}\n\n{record['input']}"
    input_str = truncate(input_str)

    return Sample(id=record["id"], input=input_str, target=record["answer"][0])


def record_to_sample_multi_element_answer(record: dict[str, Any]) -> Sample:
    """
    Convert a record from the dataset to a sample for tasks with a multi-element answer.

    Args:
        record (dict): The record from the dataset to convert to a sample.

    Returns:
        Sample: The sample with the corresponding id, input, and target.
    """
    # Construct the input string from context and input
    input_str = f"{record['context']}\n\n{record['input']}"
    input_str = truncate(input_str)

    return Sample(id=record["id"], input=input_str, target=record["answer"])


def record_to_sample_multiple_choice(record: dict[str, Any]) -> Sample:
    """
    Convert a record from the dataset to a sample for multiple choice tasks.

    Args:
        record (dict): The record from the dataset to convert to a sample.

    Returns:
        Sample: The sample with the corresponding id, input, choices, and target.
    """
    # Construct the input string from context and input
    input_str = f"{record['context']}\n\n{record['input']}"
    input_str = truncate(input_str)

    # For multiple choice tasks, the answer is always a list with a single element
    answer = record["answer"][0]

    # Convert the answer to a single letter (A, B, C, or D) for multiple choice tasks
    # to match the multiple choice solver and prompts in InfiniteBench
    index = record["options"].index(answer)
    target = chr(65 + index)  # 'A', 'B', 'C', or 'D'

    return Sample(
        id=record["id"], input=input_str, choices=record["options"], target=target
    )


def record_to_sample_code_run(record: dict[str, Any]) -> Sample:
    """
    Convert a record from the dataset to a sample for the Code.Run task.

    Args:
        record (dict): The record from the dataset to convert to a sample.

    Returns:
        Sample: The sample with the corresponding id, input, and target.
    """
    # Construct the input string from context and input
    match = re.search(r"func_\d+(?=\()", record["input"])
    if not match:
        raise ValueError(f"Could not find function name in input: {record['input']}")
    function_name = match.group(0)
    prompt = f"Following is a set of Python functions. There is a function named {function_name}."
    input_str = f"{prompt}\n\n{record['context']}\n\n{record['input']}"
    input_str = truncate(input_str)

    return Sample(id=record["id"], input=input_str, target=record["answer"][0])


def record_to_sample_math_find(record: dict[str, Any]) -> Sample:
    """
    Convert a record from the dataset to a sample for the Math.Find task.

    Args:
        record (dict): The record from the dataset to convert to a sample.

    Returns:
        Sample: The sample with the corresponding id, input, and target.
    """
    # Construct the input string from context and input
    prompt = MATH_FIND_INPUT_TO_PROMPT[record["input"]]
    input_str = f"{prompt}\n\n{record['context']}\n\n{record['input']}"
    input_str = truncate(input_str)

    return Sample(id=record["id"], input=input_str, target=record["answer"][0])
