import logging
from typing import Dict, List, Literal, NamedTuple

import numpy as np
import pandas as pd
from datasets import load_dataset  # type: ignore
from tiktoken.core import Encoding

from inspect_evals.niah.utils.text_utils import (
    check_model_context_length,
    find_nearest_period_indices,
    get_context_length_in_tokens,
    get_max_tokens,
    get_prompt_tokens,
)

logger = logging.getLogger(__name__)


class ExpParams(NamedTuple):
    """
    Holds parameters for configuring an experiment.

    Attributes:
    min_context (int): Minimum target context length.
    max_context (int): Maximum target context length.
    n_contexts (int): Number of different context lengths to generate.
    n_positions (int): Number of needle positions within each context.
    start_buffer (int): Minimum starting token index for inserting needles.
    end_buffer (int): Maximum ending token index for inserting needles.
    n_needles (int): Total number of needles to sample for insertion.
    sample_method (Literal['fixed', 'sequential', 'random']): Method for sampling needles ('fixed', 'sequential', 'random').
    fixed_index (int): Starting index for sampling if 'fixed' or 'sequential' method is used.
    n_runs (int): Number of runs for each parameter combination (of context size and position).
    main_prompt (str): Main prompt text for presenting context and task to model.
    question_prompt (str): Question prompt to be presented to model.
    token_buffer (int): Buffer token to minimise risk of exceeding maximum context.
    model_name (str): Name of the model to be used.
    """

    # Context and experimental setup
    min_context: int
    max_context: int
    n_contexts: int
    n_positions: int
    start_buffer: int
    end_buffer: int
    n_needles: int

    # Sampling setup
    sample_method: Literal["fixed", "sequential", "random"]
    fixed_index: int
    n_runs: int

    # Prompts and model setup
    main_prompt: str
    question_prompt: str
    token_buffer: int
    model_name: str


class NIAH_Data(NamedTuple):
    """
    Contains data extracted from OpenCompass NIAH dataset.

    Attributes:
    - haystacks (pd.DataFrame): DataFrame containing the various long context texts (haystacks) in which needles will be embedded.
    - needles (pd.DataFrame): DataFrame containing specific factual content (needles) with associated questions and answers.
    """

    haystacks: pd.DataFrame
    needles: pd.DataFrame


class Haystack(NamedTuple):
    """
    Contains key information for the haystack for a given NIAH experiment.

    Attributes:
    - full_context (str): Full text context of the haystack (based on max_context).
    - encoded_context (List[int]): List of tokens for full context.
    - actual_context_lengths (List[int]): List of actual context lengths to be evaluated.
    - prompt_tokens (int): Number of tokens in the prompt texts.
    """

    full_context: str
    encoded_context: List[int]
    actual_context_lengths: List[int]
    prompt_tokens: int


def read_from_hf_dataset() -> Dict[str, pd.DataFrame]:
    """
    Reads in data from the OpenCompass 'NeedleBench' Hugging Face dataset and loads them into a dictionary of pandas DataFrames.

    Subsets loaded:
        - 'en_haystack_texts': The primary context documents.
        - 'retrieval_needles': Specific needle phrases to be searched.

    Returns:
        Dict[str, pd.DataFrame]: A dictionary with subset names as keys and corresponding DataFrames as values.
    """
    # List of the haystack and needle config names
    subsets = ["en_haystack_texts", "retrieval_needles"]
    datasets: Dict[str, pd.DataFrame] = {}

    for subset in subsets:
        try:
            # Attempt to load the dataset
            dataset = load_dataset("opencompass/NeedleBench", subset)
            if "test" not in dataset:
                raise ValueError(f"'test' not found in subset: {subset}")

            # Convert to pandas DataFrame and store in the dictionary
            datasets[subset] = dataset["test"].to_pandas()

        except Exception as e:
            # Raise a RuntimeError if the dataset could not be loaded
            raise RuntimeError(f"Could not load dataset '{subset}': {str(e)}")

    return datasets


def get_data_from_hf() -> NIAH_Data:
    """
    Processes data that has been read from the OpenCompass 'NeedleBench' Hugging Face dataset and creates NIAH_Dataset object.

    Returns:
        NIAH_Data: A NameTuple containing data extracted from OpenCompass NIAH dataset.
    """
    # Get haystacks and needles dfs from HF dataset
    hf_dataset = read_from_hf_dataset()
    haystacks = hf_dataset["en_haystack_texts"]
    needles = hf_dataset["retrieval_needles"]

    # Filter and rename columns for the needle dataset
    english_needles = needles[needles["language"] == "English"].copy()
    english_needles.rename(
        columns={"retrieval_question": "question", "gold_standard_answer": "answer"},
        inplace=True,
    )

    return NIAH_Data(haystacks, english_needles)


def generate_context_lengths(
    min_context: int, max_context: int, n_contexts: int, needle_and_prompt_tokens: int
) -> List[int]:
    """
    Generates a list of linearly spaced context lengths between `min_context` and `max_context` (inclusive).

    Rounding applied to ensure context length values are integers.
    If the minimum and maximum contexts are identical, it returns a list with a single value.
    The number of contexts in list is defined by `n_contexts`.

    Args:
        min_context (int): The minimum context length.
        max_context (int): The maximum context length.
        n_contexts (int): The number of context lengths to generate.
        needle_and_prompt_tokens (int): The maximum number of tokens that the needle and prompts can contribute.

    Returns:
        List[int]: A list of rounded context lengths between `min_context` and `max_context`.
    """
    # Raise error if n_contexts < 1.
    if n_contexts < 1:
        raise ValueError("Value for `n_contexts` is less than 1.")

    # Use `max_context` if `n_contexts` == 1
    if n_contexts == 1:
        logger.warning("Number of contexts is 1, using `max_context` only.")
        return [max_context]

    # Raise error if either min_context or max_context < 0.
    if min_context < 1:
        raise ValueError("Value for `min_context` is negative.")

    if max_context < 1:
        raise ValueError("Value for `max_context` is negative.")

    # Raise error if max needle length + prompt length exceeds the min_context:
    if min_context < needle_and_prompt_tokens:
        raise ValueError(
            f"Value for `min_context` ({min_context}) is smaller the number of tokens required for needle and prompts ({needle_and_prompt_tokens})."
        )

    # Raises error if max_context is smaller than min_context
    if max_context < min_context:
        raise ValueError(
            "`max_context` is smaller then `min_context`. Cannot create context range."
        )

    # Adjust min and max context for additional tokens
    adj_min_context = min_context - needle_and_prompt_tokens
    adj_max_context = max_context - needle_and_prompt_tokens

    # Check if min_context and max_context are identical
    if adj_min_context == adj_max_context:
        logger.warning(
            "min_context and max_context are identical. Returning a list with a single value."
        )
        return [adj_min_context]

    step_size = (adj_max_context - adj_min_context) / (n_contexts - 1)
    context_lengths = [adj_min_context + i * step_size for i in range(n_contexts)]
    rounded_context_lengths = [
        round(context_length) for context_length in context_lengths
    ]

    # Check if the rounded values contain duplicates
    if len(set(rounded_context_lengths)) < n_contexts:
        logger.warning(
            "Context lengths are not unique, removing duplicates. Adjust min/max or increase n_contexts."
        )

    rounded_context_lengths = sorted(set(rounded_context_lengths))

    return rounded_context_lengths


def get_target_context_lengths(
    min_context: int, max_context: int, n_contexts: int
) -> List[int]:
    """
    Generates a list of the target context lengths based on the input parameters.

    These are useful for comparing across independent runs of the benchmark where
    the actual context lengths may vary due to differences in the sampled haystacks.

    Args:
        min_context (int): The minimum context size.
        max_context (int): The maximum context size.
        n_contexts (int): The number of contexts to be generated.

    Returns:
        List[int]: A list of target context lengths.
    """
    # Calculate the target contexts based on input parameters
    target_context_lengths = [
        int(x) for x in np.linspace(min_context, max_context, n_contexts)
    ]

    # Give warning if duplicates found in target contexts
    if len(target_context_lengths) != len(set(target_context_lengths)):
        logger.warning(
            "Duplicates found in target context lengths. Considering increasing the context range (`max_context` - `min_context`) or decreasing `n_contexts`"
        )

    return target_context_lengths


def read_and_trim_context(
    context_lengths: List[int], haystacks: pd.DataFrame, enc: Encoding
) -> List[int]:
    """
    Reads and concatenates text from a DataFrame to build a context large enough for the maximum specified length, ending on a sentence boundary.

    Args:
        context_lengths (List[int]): A list of possible context lengths, in tokens, to be considered.
        haystacks (pd.DataFrame): A DataFrame containing text snippets for context generation.
        enc (Encoding): The tokeniser encoding object, used to convert text to tokens.

    Returns:
        List[int]: The tokens for the encoded full context text, after it has been trimmed to the appropriate length
        to end on a sentence boundary.
    """
    context = ""

    # Calculate maximum context length required
    max_context_length = max(context_lengths)

    # Go through haystacks until the context is over the max length in tokens
    for text in haystacks["text"].values:
        if get_context_length_in_tokens(context, enc) >= max_context_length:
            break  # Stop once context length exceeds or meets the max length
        context += text

    # Tokenise the context
    tokens = enc.encode(context)

    # Ensure context ends at the nearest period (sentence boundary)
    period_token = enc.encode(".")  # Token for '.'

    # Define idx for slice point
    slice_idx = max_context_length

    # Forward track to find the nearest period past the max length
    while slice_idx < len(tokens) and tokens[slice_idx] not in period_token:
        slice_idx += 1

    # Return tokens for full context
    return tokens[:slice_idx]


def generate_full_context(
    hf_data: NIAH_Data, exp_params: ExpParams, enc: Encoding
) -> Haystack:
    """
    Orchestration function to create a Haystack object.

    The Haystack object contains the full context (haystack) for the longest required context
    length. This context is trimmed for shorter context evaluations in the experimental run.
    The actual context lengths to be evaluated (based on the nearest periods that occur prior
    to each of the target context lengths) are stored as well as the number of tokens present
    in the prompt texts.

    Args:
        hf_data (NIAH_Data): A NamedTuple containing data extracted from OpenCompass NIAH dataset.
        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:
        Haystack: Key information about the haystack for the current experiment.
    """
    # Get the maximum combined length for needles, questions and answers.
    max_combined_needle_tokens = get_max_tokens(
        hf_data.needles, ["needle", "question"], enc
    )

    # Get the number of tokens in the main prompt and question prompt
    prompt_tokens = get_prompt_tokens(
        exp_params.main_prompt, exp_params.question_prompt, enc
    )

    # Get the max number of tokens in the possible answers
    max_answer_tokens = get_max_tokens(hf_data.needles, ["answer"], enc)

    # Calculate number of additional tokens required
    additional_tokens = max_answer_tokens + exp_params.token_buffer

    # Calculate number of tokens in needle and prompt
    needle_and_prompt_tokens = max_combined_needle_tokens + prompt_tokens

    # Generate context lengths
    context_lengths = generate_context_lengths(
        exp_params.min_context,
        exp_params.max_context,
        exp_params.n_contexts,
        needle_and_prompt_tokens,
    )

    # Read in the context strings and create the encoded main context
    encoded_full_context = read_and_trim_context(
        context_lengths, hf_data.haystacks, enc
    )

    # Decode the tokens back into text after ensuring it ends with a period
    full_context = enc.decode(encoded_full_context)

    # Define the period token
    period_token = enc.encode(".")

    # Get nearest periods for all contexts
    actual_context_lengths = find_nearest_period_indices(
        encoded_full_context, context_lengths, period_token
    )

    # Check whether context lengths are appropriate for specified model
    check_model_context_length(
        exp_params.model_name, actual_context_lengths, additional_tokens
    )

    # Return Haystack object
    return Haystack(
        full_context, encoded_full_context, actual_context_lengths, prompt_tokens
    )
