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

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

logger = logging.getLogger(__name__)


def sample_needles(
    needle_df: pd.DataFrame,
    n_samples: int,
    sample_method: Literal["fixed", "sequential", "random"] = "random",
    fixed_index: int = 0,
) -> pd.DataFrame:
    """
    Samples `n_samples` rows from the DataFrame using the specified method.

    Args:
        needle_df (pd.DataFrame): The DataFrame containing needle data to sample from.
        n_samples (int): The number of samples to return.
        sample_method (Literal['fixed', 'sequential', 'random']): The sampling method; can be 'fixed', 'sequential', or 'random'.
            - 'fixed': Selects a single needle multiple times from a fixed position (`fixed_index`).
            - 'sequential': Starts sampling from `fixed_index` and proceeds sequentially, wrapping around if needed.
            - 'random': Randomly samples needles, allowing repeats if `n_samples` exceeds available rows.
        fixed_index (int): The starting index for 'fixed' and 'sequential' sampling.

    Returns:
        pd.DataFrame: A DataFrame containing the sampled needles.
    """
    rows = len(needle_df)

    if rows == 0:
        raise ValueError("DataFrame is empty.")

    # Fixed sampling: always the same row.
    if sample_method == "fixed":
        if fixed_index < 0 or fixed_index >= rows:
            raise ValueError(
                "Invalid fixed_index. Must be between 0 and number of rows - 1."
            )
        return needle_df.iloc[[fixed_index] * n_samples]

    # Sequential sampling: select rows sequentially, wrap around if necessary.
    elif sample_method == "sequential":
        indices = [(fixed_index + i) % rows for i in range(n_samples)]
        return needle_df.iloc[indices]

    # Random sampling: allow repeats if n_samples > rows.
    elif sample_method == "random":
        if n_samples > rows:
            # Repeat the rows until we have enough samples, and then add any extra samples
            random_indices = random.sample(range(rows), rows)
            full_repeats = n_samples // rows
            extra_samples = n_samples % rows
            random_indices = (random_indices * full_repeats) + random.sample(
                range(rows), extra_samples
            )
        else:
            random_indices = random.sample(range(rows), n_samples)

        return needle_df.iloc[random_indices]

    else:
        raise ValueError(
            "Invalid sample method. Use 'fixed', 'sequential', or 'random'."
        )


def stretch_or_squash_needles(
    sampled_needles: pd.DataFrame, n_combos: int
) -> pd.DataFrame:
    """
    Adjusts the DataFrame containing sampled needles to match the required number by either truncating or repeating rows.

    Args:
        sampled_needles (pd.DataFrame): A DataFrame containing the sampled needle data.
        n_combos (int): The required number of needles.

    Returns:
        pd.DataFrame: A DataFrame containing exactly `n_combos` rows by either truncating or repeating
        rows from `sampled_needles`.
    """
    # If sampled needles exceeds or matches required number of needles, squash to correct size.
    n_samples = len(sampled_needles)
    if n_samples >= n_combos:
        return sampled_needles.head(n_combos)

    # Otherwise stretch to correct size.
    repetitions = (n_combos // n_samples) + 1
    stretched_sample = pd.concat(
        [sampled_needles] * repetitions, ignore_index=True
    ).head(n_combos)

    return stretched_sample


def repeat_with_shift(df: pd.DataFrame, n_reps: int) -> pd.DataFrame:
    """
    Repeats the rows of a DataFrame `n_reps` times, shifting the order by one for each repetition.

    Args:
        df (pd.DataFrame): The DataFrame whose rows are to be repeated and shifted.
        n_reps (int): The number of times to repeat the DataFrame with shifts.

    Returns:
        pd.DataFrame: A DataFrame with repeated rows, shifted by one row on each repetition.
    """
    rows = []
    num_rows = len(df)

    # Create the repeated rows with shifting
    for rep_idx in range(n_reps):
        for row_idx in range(num_rows):
            # Append the row index with the correct shifting logic
            shifted_index = (row_idx + rep_idx) % num_rows
            rows.append(df.iloc[shifted_index].values)  # Get the row values

    # Create a new DataFrame from the rows list
    final_df = pd.DataFrame(rows, columns=df.columns)

    # Reset the index to be continuous at the end
    final_df.reset_index(drop=True, inplace=True)

    return final_df


def insert_needle_into_context(
    context: str, needle: str, insertion_point: int, enc: Encoding
) -> str:
    """
    Inserts a needle into the context at a specified position, making sure insertion is at a sentence boundary.

    Args:
        context (str): The original context text.
        needle (str): The text to be inserted.
        insertion_point (int): The approximate token position for insertion.
        enc (Encoding): The tokeniser encoding object, used to convert text to tokens.

    Returns:
        str: The context with the needle inserted at a sentence boundary nearest to `insertion_point`.
    """
    # Encode the needle and context into tokens
    needle_tokens = enc.encode(needle)
    context_tokens = enc.encode(context)

    # Capture the tokens up to the insertion point
    tokens_before_insertion = context_tokens[:insertion_point]

    # Encode the period (".") to match the token for sentence endings
    period_token = enc.encode(".")

    # Backtrack from the insertion point to find the nearest sentence-ending token
    while tokens_before_insertion and tokens_before_insertion[-1] not in period_token:
        insertion_point -= 1
        tokens_before_insertion = context_tokens[:insertion_point]

    # Define whitespace token
    whitespace_token = enc.encode(" ")

    # Insert the needle after the identified sentence end
    new_tokens = (
        tokens_before_insertion
        + whitespace_token
        + needle_tokens
        + context_tokens[insertion_point:]
    )

    # Decode the final token sequence back into a string and return it
    updated_context = enc.decode(new_tokens)
    return updated_context


def generate_needle_positions(
    n_positions: int, context_length: int, start_buffer: int, end_buffer: int
) -> List[int]:
    """
    Generates a list of needle positions based on the specified context length and buffer values.

    Ensures that the needle positions fit within the bounds defined by the start and end buffers.

    Args:
        n_positions (int): The number of needle positions to generate.
        context_length (int): The total context length.
        start_buffer (int): The buffer of tokens to leave at the start of the context.
        end_buffer (int): The buffer of tokens to leave at the end of the context.

    Returns:
        List[int]: A list of rounded needle positions within the context bounds.
    """
    # Raise error when n_positions < 1.
    if n_positions < 1:
        raise ValueError("Value for `n_positions` is less than 1.")

    # Raise error if buffer values are too large for context size.
    if (
        start_buffer > context_length
        or end_buffer > context_length
        or (start_buffer + end_buffer) >= context_length
    ):
        raise ValueError(
            f"Buffer values are too large for context size: {context_length}."
        )

    start_token = start_buffer
    end_token = context_length - end_buffer

    # Handle the case when n_positions == 1 by selecting position in the middle of start and end buffers.
    if n_positions == 1:
        rounded_needle_positions = [round((start_token + end_token) / 2)]
    else:
        step_size = (end_token - start_token) / (n_positions - 1)
        needle_positions = [start_token + i * step_size for i in range(n_positions)]
        rounded_needle_positions = [
            round(needle_position) for needle_position in needle_positions
        ]

    # Check if the rounded values contain duplicates
    if len(set(rounded_needle_positions)) < n_positions:
        logger.warning(
            f"Needle positions are not unique, removing duplicates. Reduce top/bottom buffers or increase context length: {context_length}"
        )
        rounded_needle_positions = list(sorted(set(rounded_needle_positions)))

    return rounded_needle_positions


def generate_needle_positions_across_contexts(
    n_positions: int, context_lengths: List[int], start_buffer: int, end_buffer: int
) -> Dict[int, List[int]]:
    """
    Generates a exact needle positions (in tokens) for each context length.

    Args:
        n_positions (int): The number of needle positions for each context length.
        context_lengths (List[int]): List of context lengths to generate needle positions for.
        start_buffer (int): The buffer zone (in tokens) from the start of the context, within which positions won't be generated.
        end_buffer (int): The buffer zone (in tokens) from the end of the context, within which positions won't be generated.

    Returns:
        Dict[int, List[int]]: A dictionary where keys are context lengths and values are lists of positions
        within those context lengths.
    """
    needle_positions = {}

    for context_length in context_lengths:
        needle_positions[context_length] = generate_needle_positions(
            n_positions, context_length, start_buffer, end_buffer
        )

    return needle_positions


def get_target_needle_position_pencentiles(n_positions: int) -> List[int]:
    """
    Calculates the target needle positions as percentiles.

    Required for aggregating results.
    Note: Percentiles are relative to the insertable region of the context (between the start and end buffers).

    Args:
        n_positions (int): The number of needle positions required.

    Returns:
        List(int): The target needle positions as percentiles.
    """
    return [int(x) for x in np.linspace(0, 100, n_positions)]
