from typing import Any, Dict, List

import pandas as pd
from inspect_ai.dataset import MemoryDataset, Sample
from tiktoken.core import Encoding

from inspect_evals.niah.utils.dataset_generation import (
    ExpParams,
    Haystack,
    NIAH_Data,
    get_target_context_lengths,
)
from inspect_evals.niah.utils.needle_utils import (
    generate_needle_positions_across_contexts,
    get_target_needle_position_pencentiles,
    insert_needle_into_context,
    repeat_with_shift,
    sample_needles,
    stretch_or_squash_needles,
)
from inspect_evals.niah.utils.text_utils import find_nearest_period_indices


def record_to_sample(record: dict[str, Any]) -> Sample:
    """
    Converts a dictionary record into a `Sample` object by extracting necessary fields for input, target, and metadata.

    Args:
        record (Dict[str, Any]): A dictionary containing the data for the sample.
                                 The dictionary should include the following keys:
                                 - "text" (str): The context with needle.
                                 - "question" (str): The question to be asked about the needle.
                                 - "answer" (str): The correct answer to the question.
                                 - "context_length_target" (Optional): The target context length for the given run.
                                 - "position_target" (Optional): The target position for the given run.
                                 - "context_length_actual" (Optional): The actual context length used for the given run.
                                 - "position_actual" (Optional): The actual position used for the given run.

    Returns:
        Sample: A `Sample` object containing:
                - `input`: A generated prompt including the "text" and "question" from the record.
                - `target`: The correct answer.
                - `metadata`: Experimental parameters (context length and position).
    """
    return Sample(
        input=record["needle"],
        target=record["answer"],
        metadata={
            # Including key experimental params for inserting into haystack and subsetting
            "haystack_length": record["haystack_length"],
            "full_context_length": record["full_context_length"],
            "target_context_length": record["target_context_length"],
            "position": record["position"],
            "target_position": record["target_position"],
            "needle_question": record["question"],
        },
    )


def sample_df_to_dataset(samples_df: pd.DataFrame) -> MemoryDataset:
    """
    Converts DataFrame to MemoryDataset by transforming rows into a corresponding sample using the `record_to_sample` function.

    Args:
        samples_df (pd.DataFrame): The input DataFrame containing the data to be converted. Each row in
                           the DataFrame represents a record that will be transformed into a sample.

    Returns:
        MemoryDataset: A dataset containing all the samples derived from the DataFrame records,
                       wrapped in a `MemoryDataset` object.
    """
    all_samples = [
        record_to_sample(record.to_dict()) for _, record in samples_df.iterrows()
    ]
    return MemoryDataset(all_samples)


def create_sample_dataset(
    context_lengths: List[int],
    sampled_needles: pd.DataFrame,
    needle_positions: Dict[int, List[int]],
    n_runs: int,
    prompt_tokens: int,
    target_context_lengths: List[int],
    target_needle_positions: List[int],
    enc: Encoding,
) -> pd.DataFrame:
    """
    Generates a dataset to create Samples based on combinations of context lengths and needle positions for given number of runs.

    Args:
        context_lengths (List[int]): Context lengths for each needle insertion.
        sampled_needles (pd.DataFrame): DataFrame containing needle data with text, questions, and answers.
        needle_positions (Dict[int, List[int]]): A dictionary of target positions for needle insertions by context length.
        n_runs (int): Number of runs for generating variations of the dataset.
        prompt_tokens (int): The additional number of tokens that the task prompt will contribute.
        target_context_lengths (List[int]): The target context lengths for each sample.
        target_needle_positions (List[int]): The target needle position for each sample.
        enc (Encoding): Tokeniser encoding used to tokenise and process text.

    Returns:
        pd.DataFrame: The generated dataset in Pandas DataFrame format, with a row for each Sample to be generated and the following columns:
            - 'context_length_target': Target length of the context.
            - 'position': Position for needle insertion.
            - 'needle': The needle text inserted.
            - 'question': The retrieval question associated with the needle.
            - 'answer': The expected answer to the question.
    """
    # Create list for storing sample values
    dataset_values = []

    # Iterate through multiple runs if required
    for run_idx in range(n_runs):
        # Iterate through context lengths
        for context_idx, context_length in enumerate(context_lengths):
            # Retrieve needle positions for given context length
            context_needle_positions = needle_positions[context_length]

            # Iterate through needle positions
            for position_idx, needle_pos in enumerate(context_needle_positions):
                # Get needle index and select from sampled needles
                needle_idx = (
                    position_idx
                    + (context_idx * len(context_needle_positions))
                    + (run_idx * len(context_needle_positions) * len(context_lengths))
                )
                needle_row = sampled_needles.iloc[needle_idx]

                needle_tokens = len(enc.encode(needle_row["needle"]))
                question_tokens = len(enc.encode(needle_row["question"]))

                # Store the sample values as list of dicts
                dataset_values.append(
                    {
                        "haystack_length": context_length,
                        "full_context_length": context_length
                        + prompt_tokens
                        + needle_tokens
                        + question_tokens,
                        "target_context_length": target_context_lengths[context_idx],
                        "position": needle_pos,
                        "target_position": target_needle_positions[position_idx],
                        "needle": needle_row["needle"],
                        "question": needle_row["question"],
                        "answer": needle_row["answer"],
                    }
                )
    # Return as dataframe
    return pd.DataFrame(dataset_values)


def needle_into_haystack(
    encoded_full_context: List[int],
    needle: str,
    context_length: int,
    needle_position: int,
    enc: Encoding,
) -> str:
    """
    Generates a string with a given target fact for retrieval ('needle') inserted into a larger context of specified size ('haystack').

    Args:
        encoded_full_context (List[int]): List of tokens for encoded complete context text to be trimmed to specified size.
        needle (str): The fact ('needle') to be inserted into context ('haystack').
        context_length (int): The required context length for current context/haystack.
        needle_position (int): The target (ideal) token index for the needle to be inserted into the context.
        enc (Encoding): Tokeniser encoding used to tokenise and process text.

    Returns:
        str: The context of specified length with needle insert at the specified position.
    """
    # Define the period token
    period_token = enc.encode(".")

    # Get required length
    actual_context_length = find_nearest_period_indices(
        encoded_full_context, [context_length], period_token
    )[0]

    # Slice the encoded context to required length
    encoded_selected_context = encoded_full_context[:actual_context_length]

    # Find the nearest period for the given needle position
    actual_position = find_nearest_period_indices(
        encoded_selected_context, [needle_position], period_token
    )[0]

    # Decoded the trimmed context
    selected_context = enc.decode(encoded_selected_context)

    # Insert needle
    text_with_needle = insert_needle_into_context(
        selected_context, needle, actual_position, enc
    )

    return text_with_needle


def generate_samples(
    hf_data: NIAH_Data, haystack: Haystack, exp_params: ExpParams, enc: Encoding
) -> pd.DataFrame:
    """
    Generates a dataset of samples.

    This function orchestrates the generation of needle positions, sampling needles,
    and creating experimental datasets. It considers context lengths, needle positions,
    and sampling methods specified in the experimental parameters.

    Args:
        hf_data (NIAH_Data): A NamedTuple containing data extracted from OpenCompass NIAH dataset.
        haystack (Haystack): Haystack object containing complete context (haystack) in which the needle will be embedded.
        exp_params (ExpParams): A NamedTuple containing the experimental parameters for the evaluation.
        enc (Encoding): The tokeniser encoding object, used to convert text to tokens.

    Returns:
        pd.DataFrame: A pandas DataFrame containing the generated dataset of samples
        including context lengths, needle positions, and other required information.
    """
    # Generate needle positions
    needle_positions = generate_needle_positions_across_contexts(
        exp_params.n_positions,
        haystack.actual_context_lengths,
        exp_params.start_buffer,
        exp_params.end_buffer,
    )

    # Update n_positions and n_contexts if adjusted during context and needle position generation
    n_contexts = len(haystack.actual_context_lengths)
    n_positions = len(next(iter(needle_positions.values())))

    target_context_lengths = get_target_context_lengths(
        exp_params.min_context, exp_params.max_context, n_contexts
    )
    target_needle_positions = get_target_needle_position_pencentiles(n_positions)

    # Calculate number of experimental combinations
    n_combos = n_positions * n_contexts

    # Sample needles based on the sampling method
    sampled_needles = sample_needles(
        hf_data.needles,
        exp_params.n_needles,
        exp_params.sample_method,
        exp_params.fixed_index,
    )

    # Get the right number of needles for the number of combinations
    experimental_needles = stretch_or_squash_needles(sampled_needles, n_combos)

    # Repeat and shift for multiple runs
    all_needles = repeat_with_shift(experimental_needles, exp_params.n_runs)

    # Create the dataset of samples
    sample_dataset = create_sample_dataset(
        haystack.actual_context_lengths,
        all_needles,
        needle_positions,
        exp_params.n_runs,
        haystack.prompt_tokens,
        target_context_lengths,
        target_needle_positions,
        enc,
    )

    return sample_dataset
