from typing import Literal

import tiktoken
from inspect_ai import Task, task
from inspect_ai.model import get_model
from inspect_ai.solver import (
    Generate,
    Solver,
    TaskState,
    generate,
    prompt_template,
    solver,
)
from tiktoken.core import Encoding

from inspect_evals.niah.utils.dataset_generation import (
    ExpParams,
    Haystack,
    generate_full_context,
    get_data_from_hf,
)
from inspect_evals.niah.utils.prompting import (
    MAIN_PROMPT,
    QUESTION_PROMPT,
    create_question_prompt,
)
from inspect_evals.niah.utils.sample_generation import (
    generate_samples,
    needle_into_haystack,
    sample_df_to_dataset,
)
from inspect_evals.niah.utils.scoring import (
    custom_model_graded_qa_with_history_scorer,
    return_metadata_variable_as_history,
)
from inspect_evals.niah.utils.text_utils import get_model_or_default

# Define a token buffer for max context to avoid potential edge case issues with long outputs
TOKEN_BUFFER = 100


@task
def niah(
    min_context: int = 10000,
    max_context: int = 120000,
    n_contexts: int = 15,
    n_positions: int = 15,
    start_buffer: int = 0,
    end_buffer: int = 0,
    n_needles: int = 1,
    sample_method: Literal["fixed", "sequential", "random"] = "fixed",
    fixed_index: int = 0,
    n_runs: int = 1,
) -> Task:
    """
    Inspect Task implementation for NIAH (Needle in a Haystack).

    This function generates a task that evaluates the model on a dataset with varying context sizes and needle positions.
    Needles are inserted into the context to measure the model's ability to retrieve relevant information.

    Args:
        min_context (int): Minimum context length to evaluate. Default is 10000.
        max_context (int): Maximum context length to evaluate. Default is 120000.
        n_contexts (int): The number of contexts to evaluate. Default is 15.
        n_positions (int): The number of positions to evaluate for a given context length. Default is 15.
        start_buffer (int): Buffer at the top of the context to avoid placing needles. Default is 0.
        end_buffer (int): Buffer at the bottom of the context to avoid placing needles. Default is 0.
        n_needles (int): The number of needles to sample. Default is 1.
        sample_method (Literal["fixed", "random"]): Method for sampling the needles.
            If "fixed", a single specific needle index is used for all trials.
            If "random", a new needle is randomly sampled for each trial.
            If "sequential", needles are sequentially sampled across trials.
            Default is "fixed".
        fixed_index (int): The index of the needle to use when `sample_method` is "fixed" or
            the index of the starting position when `sample_method` is "sequential". Default is 0.
        n_runs (int): The number of runs per set of experimental parameters. Default is 1.

    Returns:
        Task: A Task object containing the dataset, the solver configuration, and a custom scorer with metadata handling.
    """
    # Get the active model
    model = get_model()

    # Use default model name for tokenisation if no tokeniser found for current model
    tokeniser_model_name = get_model_or_default(model.name)

    # Create an encoder for given model
    enc = tiktoken.encoding_for_model(tokeniser_model_name)

    # Import OpenCompass 'Needle in a Haystack' dataset from HF
    hf_data = get_data_from_hf()

    # Generate ExpParams object for storing experimental parameters.
    exp_params = ExpParams(
        min_context=min_context,
        max_context=max_context,
        n_contexts=n_contexts,
        n_positions=n_positions,
        start_buffer=start_buffer,
        end_buffer=end_buffer,
        n_needles=n_needles,
        sample_method=sample_method,
        fixed_index=fixed_index,
        n_runs=n_runs,
        main_prompt=MAIN_PROMPT,
        question_prompt=QUESTION_PROMPT,
        token_buffer=TOKEN_BUFFER,
        model_name=model.name,
    )

    # Generate the haystack for the largest context length required. Smaller context lengths will trim haystack.
    haystack = generate_full_context(hf_data, exp_params, enc)

    # Generate a DataFrame with Sample information
    samples_df = generate_samples(hf_data, haystack, exp_params, enc)

    # Convert Sample DataFrame to Dataset
    dataset = sample_df_to_dataset(samples_df)

    # Return the Task
    return Task(
        dataset=dataset,
        solver=[
            add_to_haystack(
                haystack, enc
            ),  # Take needle and other information from Sample to generate combined haystack and needle text.
            prompt_template(MAIN_PROMPT),
            generate(),
        ],
        # Custom wrapper used to allow for grouped scoring and parsing metadata to scorer
        scorer=custom_model_graded_qa_with_history_scorer(
            include_history=return_metadata_variable_as_history,
        ),
    )


@solver
def add_to_haystack(haystack: Haystack, enc: Encoding) -> Solver:
    """
    Custom solver function.

    Inserts a specified prompt (needle) into a larger context (haystack) string based on provided Sample metadata parameters

    Args:
        haystack (Haystack): Haystack object containing complete context (haystack) in which the needle will be embedded.
        enc (Encoding): The tokeniser encoding object, used to convert text to tokens.

    Returns:
        Solver: An asynchronous solver function that takes `TaskState` and `Generate`
        and returns an updated `TaskState` with prompt text.
    """

    async def solve(state: TaskState, generate: Generate) -> TaskState:
        prompt = state.user_prompt
        metadata = state.metadata
        full_context = needle_into_haystack(
            haystack.encoded_context,
            prompt.text,
            metadata["haystack_length"],
            metadata["position"],
            enc,
        )
        prompt.text = create_question_prompt(full_context, metadata["needle_question"])
        return state

    return solve
