import logging
from typing import Any, List

import pandas as pd
import tiktoken
from tiktoken.core import Encoding

from inspect_evals.niah.utils.model_contexts import MODEL_CONTEXT_LENGTHS

logger = logging.getLogger(__name__)


def get_context_length_in_tokens(context: str, enc: Encoding) -> int:
    """
    Calculates the tokenised length of a given text context.

    Args:
        context (str): The context whose length is to be measured in tokens.
        enc (Encoding): The tokeniser encoding object, used to convert text to tokens.

    Returns:
        int: The number of tokens in the given context.
    """
    return len(enc.encode(context))


def get_max_tokens(needles: pd.DataFrame, columns: List[str], enc: Encoding) -> int:
    """
    Calculates the maximum token count across specific columns in a DataFrame of needle texts.

    Args:
        needles (pd.DataFrame): The DataFrame containing needle data.
        columns (List[str]): List of column names in `needles` DataFrame for which the token count is calculated.
        enc (Encoding): The tokeniser encoding object, used to convert text to tokens.

    Returns:
        int: The maximum combined token count for any single row across the specified columns.
    """

    # Function to get token count for a single row across relevant columns
    def combined_token_count(row: Any) -> int:
        """
        Calculates the total token count for a single row across relevant columns.

        Args:
            row (pd.Series): A row in the DataFrame.

        Returns:
            int: The total token count for the values in the row.
        """
        total_tokens = 0

        # Tokenise the content in the value for each row and count tokens
        for value in row:
            total_tokens += len(enc.encode(str(value)))
        return total_tokens

    # Apply the function to each row for relevant columns then find max
    max_combined_tokens = needles[columns].apply(combined_token_count, axis=1).max()

    return max_combined_tokens


def get_prompt_tokens(main_prompt: str, question_prompt: str, enc: Encoding) -> int:
    """
    Calculates the total token count for the combined main prompt and question prompt.

    Args:
        main_prompt (str): The main prompt text.
        question_prompt (str): The question prompt text.
        enc (Encoding): The tokeniser encoding object, used to convert text to tokens.

    Returns:
        int: The combined token count of the main prompt and question prompt.
    """
    return len(enc.encode(str(main_prompt))) + len(enc.encode(str(question_prompt)))


def get_model_or_default(model_name: str) -> str:
    """
    Check if the provided model name is a recognised OpenAI model.

    If the model name is recognised, return it; otherwise, return a default OpenAI model ("gpt-3.5-turbo").

    Args:
        model_name (str): The model name to check.

    Returns:
        str: The recognised model name or default model name (if it's not recognised).
    """
    default_model = "gpt-3.5-turbo"

    # Check if the model name is recognised
    try:
        tiktoken.encoding_for_model(model_name)
    except Exception as e:
        # Default to a specific model if the provided name is not recognised
        logger.warning(f"Issue finding tokeniser for {model_name}: '{e}'.")
        logger.warning(f"Using {default_model} tokeniser as default.")
        return default_model

    # Return model name if it's a valid model name for tiktoken
    return model_name


def check_model_context_length(
    model_name: str, context_lengths: List[int], additional_tokens: int
) -> None:
    """
    Checks if the maximum combined context length exceeds the context length for a given model.

    If the context length is not known for the model, it raises a warning.
    If the context length exceeds the limit, it raises a ValueError.

    Args:
        model_name (str): The name of the model to check context length for.
        context_lengths (List[int]): A list of context lengths.
        additional_tokens (int): Additional tokens to be accounted for.

    Raises:
        ValueError: If the context length exceeds the model's context length.
        UserWarning: If the context length limit for the model is unknown.
    """
    if model_name in MODEL_CONTEXT_LENGTHS:
        if MODEL_CONTEXT_LENGTHS[model_name] < max(context_lengths) + additional_tokens:
            exceed_amount = (
                max(context_lengths)
                + additional_tokens
                - MODEL_CONTEXT_LENGTHS[model_name]
            )
            raise ValueError(
                f"Specified context lengths could exceed context length limit for model. Reduce by {exceed_amount} tokens."
            )
    else:
        logger.warning(
            f"Context length for model '{model_name}' not found. Unable to check whether specified context lengths are within limits."
        )


def find_nearest_period_indices(
    encoded_full_context: List[int], context_lengths: List[int], period_token: List[int]
) -> List[int]:
    """
    Determines the index of the nearest period (sentence boundary) for each target length in `context_lengths`.

    Helps ensure that contexts are truncated at sentence boundaries within the specified token lengths.

    Args:
        encoded_full_context (List[int]): List of tokens for encoded complete context text from which to find period indices.
        context_lengths (List[int]): Target context lengths for which nearest period indices are calculated.
        period_token (List[int]): Tokenisation of the `.` (period) character.

    Returns:
        List[int]: A list of indices, each representing the position in tokens where the nearest period
        occurs for each target length in `context_lengths`.
    """
    period_indices = []

    for context_length in context_lengths:
        # Check that context length of `full_context` is large enough
        if len(encoded_full_context) < context_length:
            raise ValueError(
                f"Required context length ({context_length}) is larger than the full context length ({len(encoded_full_context)})."
            )

        # If the context length is more than available tokens, adjust to the max possible
        target_tokens = encoded_full_context[:context_length]

        # Backtrack to find the nearest period to this target length
        while target_tokens and target_tokens[-1] not in period_token:
            target_tokens = target_tokens[:-1]

        # Calculate the index of the closest period
        period_index = len(target_tokens)
        period_indices.append(period_index)

    return period_indices
