import re
from typing import Literal

import tiktoken
from inspect_ai.model import get_model

from inspect_evals.infinite_bench.constants import TOKEN_BUFFER, TOKEN_TO_CHAR_RATIO
from inspect_evals.infinite_bench.utils import get_model_context_size


def find_nearest_boundary(
    text: str, start_index: int, direction: Literal["left", "right"]
) -> int:
    """
    Find the nearest whitespace or symbol boundary.

    Args:
        text (str): The text to search for the nearest boundary.
        start_index (int): The index to start searching from.
        direction (str): The direction to search for the nearest boundary.
            - "left": Search for the nearest boundary to the left.
            - "right": Search for the nearest boundary to the right.

    Returns:
        int: The index of the nearest boundary.
    """
    regex = r"[\s,;:{}()\[\]]"
    search_text = (
        text[:start_index][::-1] if direction == "left" else text[start_index:]
    )
    match = re.search(regex, search_text)
    # If no match is found, return the start_index
    if not match:
        return start_index
    return (
        start_index - match.start()
        if direction == "left"
        else start_index + match.start()
    )


def get_model_encoding(model_name: str) -> tiktoken.Encoding:
    """
    Get the encoding for the GPT model dynamically.

    - Use tiktoken with the given model name.
    - Default to using tiktoken with "gpt-4o-mini" if the model name is not recognized.

    Args:
        model_name (str): The name of the model to get the encoding for.

    Returns:
        tiktoken.Encoding: The encoding for the model.
    """
    try:
        return tiktoken.encoding_for_model(model_name)
    except Exception as e:
        print(f"Model {model_name} not recognized, defaulting to gpt-4o-mini: {e}")
        return tiktoken.encoding_for_model("gpt-4o-mini")


# Based on source code from https://github.com/OpenBMB/InfiniteBench
def truncate_by_tokens(input_text: str, model_name: str) -> str:
    """
    Truncate input based on model-specific max tokens.

    Args:
        input_text (str): The input text to be truncated.
        model_name (str): The name of the model to truncate the input for.

    Returns:
        str: The truncated input text.
    """
    max_tokens = get_model_context_size(model_name)
    max_tokens = max_tokens - TOKEN_BUFFER

    # Get the correct encoding for the model
    encoding = get_model_encoding(model_name)

    # Tokenize the input text using tiktoken
    tokens = encoding.encode(input_text)
    print(f"Before truncation: {len(tokens)}")

    # Apply truncation if the number of tokens exceeds max_tokens
    if len(tokens) <= max_tokens:
        return input_text  # No truncation needed if within max_length

    split = max_tokens // 2  # Split the sequence into two halves
    tokens = (
        tokens[:split] + tokens[-split:]
    )  # Keep the first and last parts, remove the middle
    print(f"After truncation: {len(tokens)}")

    # Ensure truncation was successful
    assert len(tokens) <= max_tokens

    # Decode the truncated tokens back to text using tiktoken
    return encoding.decode(tokens)


def truncate_raw_text(input_text: str, model_name: str) -> str:
    """
    Truncate raw text based on the model's token limit by estimating the max character length.

    Applies whitespace and symbol-based truncation around the cutoff.

    Args:
        input_text (str): The text to be truncated.
        model_name (str): The name of the model to determine the max context size.

    Returns:
        str: The truncated text.
    """
    # Get the max tokens for the model, default to 512 if not found
    max_tokens = get_model_context_size(model_name)

    # Adjust max tokens with the buffer
    effective_max_tokens = max_tokens - TOKEN_BUFFER

    # Estimate the max character length based on the token-to-character ratio
    max_length = int(effective_max_tokens * TOKEN_TO_CHAR_RATIO)

    # Print the original input length
    print(f"Original input length (characters): {len(input_text)}")
    print(f"Estimated max character length for {model_name}: {max_length}")

    if len(input_text) <= max_length:
        return input_text  # No truncation needed if within max_length

    split = max_length // 2  # Split the sequence into two halves
    truncated_text = (
        input_text[:split] + input_text[-split:]
    )  # Keep the first and last parts, remove the middle

    # Apply whitespace and symbol-based truncation at the boundaries
    # Find the nearest boundaries on both sides
    split = max_length // 2
    left_boundary = find_nearest_boundary(truncated_text, split, "left")
    right_boundary = find_nearest_boundary(truncated_text, split, "right")

    # Trim the truncated text around the nearest boundaries and add ellipses
    refined_text = (
        truncated_text[:left_boundary].rstrip()
        + "..."
        + truncated_text[right_boundary:].lstrip()
    )

    # Print the truncated input length
    print(f"Refined input length (characters): {len(refined_text)}")

    return refined_text


def truncate(input_text: str) -> str:
    """
    Truncate input text based on the model's token limit.

    - If the model is a GPT model, truncate by tokens.
    - Otherwise, truncate by raw text.

    Args:
        input_text (str): The input text to be truncated.

    Returns:
        str: The truncated input text.
    """
    model_name = get_model().name
    if model_name.startswith("gpt"):
        return truncate_by_tokens(input_text, model_name)
    else:
        return truncate_raw_text(input_text, model_name)
