from unittest.mock import MagicMock

import pandas as pd
import pytest

from inspect_evals.niah.utils import (
    dataset_generation,
    needle_utils,
    sample_generation,
    text_utils,
)


@pytest.mark.dataset_download
def test_download_niah_dataset():
    """Test that we can download and load the NIAH dataset from HuggingFace."""
    # Attempt to download and load the dataset
    dataset_dict = dataset_generation.read_from_hf_dataset()

    # Check we got the expected dataframes
    assert "en_haystack_texts" in dataset_dict
    assert "retrieval_needles" in dataset_dict

    # Check the dataframes have content
    assert len(dataset_dict["en_haystack_texts"]) > 0
    assert len(dataset_dict["retrieval_needles"]) > 0

    # Check the dataframes have the expected columns
    haystack_df = dataset_dict["en_haystack_texts"]
    needles_df = dataset_dict["retrieval_needles"]

    assert "text" in haystack_df.columns
    assert "needle" in needles_df.columns
    assert "question" in needles_df.columns
    assert "answer" in needles_df.columns


class TestTokenFuncs:
    def test_get_max_tokens(self):
        mock_needles = pd.DataFrame(
            {"needle": ["needle1", "needle2"], "question": ["question1", "question2"]}
        )

        # Simple mock function to replace encoding for testing, fetches list of chars
        mock_enc = MagicMock()
        mock_enc.encode.side_effect = lambda x: list(x)

        result = text_utils.get_max_tokens(
            mock_needles, ["needle", "question"], mock_enc
        )

        expected_max_tokens = max(
            len("needle1") + len("question1"), len("needle2") + len("question2")
        )
        assert result == expected_max_tokens

    def test_get_prompt_tokens(self):
        mock_enc = MagicMock()
        mock_enc.encode.side_effect = lambda x: list(x)

        main_prompt = "This is a main prompt."
        question_prompt = "This is a question prompt."

        result = text_utils.get_prompt_tokens(main_prompt, question_prompt, mock_enc)
        expected_tokens = len(main_prompt) + len(question_prompt)
        assert result == expected_tokens


class TestGenerateContextLengths:
    def test_valueerror_n_contexts_less_than_1(self):
        with pytest.raises(ValueError, match="Value for `n_contexts` is less than 1."):
            dataset_generation.generate_context_lengths(10, 100, 0, 5)

    def test_valueerror_min_context_negative(self):
        with pytest.raises(ValueError, match="Value for `min_context` is negative."):
            dataset_generation.generate_context_lengths(-10, 100, 5, 5)

    def test_valueerror_max_context_negative(self):
        with pytest.raises(ValueError, match="Value for `max_context` is negative."):
            dataset_generation.generate_context_lengths(10, -100, 5, 5)

    def test_valueerror_min_context_less_than_needle_and_prompt_tokens(self):
        with pytest.raises(
            ValueError,
            match=r"Value for `min_context` \(10\) is smaller the number of tokens required for needle and prompts \(15\)\.",
        ):
            dataset_generation.generate_context_lengths(10, 100, 5, 15)

    def test_valueerror_max_context_smaller_than_min_context(self):
        with pytest.raises(
            ValueError,
            match="`max_context` is smaller then `min_context`. Cannot create context range.",
        ):
            dataset_generation.generate_context_lengths(100, 50, 5, 10)

    def test_single_value_when_adj_min_context_equals_adj_max_context(self):
        result = dataset_generation.generate_context_lengths(110, 110, 5, 10)
        assert result == [100]

    def test_rounded_context_lengths_are_calculated_correctly(self):
        min_context = 150
        max_context = 351
        n_contexts = 5
        needle_and_prompt_tokens = 50

        result = dataset_generation.generate_context_lengths(
            min_context, max_context, n_contexts, needle_and_prompt_tokens
        )

        expected_result = [100, 150, 200, 251, 301]
        assert result == expected_result

    def test_removes_duplicates_from_rounded_context_lengths(self):
        min_context = 90
        max_context = 100
        n_contexts = 12
        needle_and_prompt_tokens = 0

        result = dataset_generation.generate_context_lengths(
            min_context, max_context, n_contexts, needle_and_prompt_tokens
        )

        expected_result = [90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100]
        assert result == expected_result


class TestSampleGenerationFunctions:
    def test_sample_needles_with_n_samples(self):
        # Test that the correct number of samples are returned
        needles_df = pd.DataFrame(
            {
                "needle": ["needle1", "needle2", "needle3"],
                "question": ["q1", "q2", "q3"],
                "answer": ["a1", "a2", "a3"],
            }
        )

        n_samples = 2
        result_df = needle_utils.sample_needles(needles_df, n_samples)

        expected_rows = n_samples
        assert len(result_df) == expected_rows

    def test_generate_needle_positions_across_contexts(self):
        n_positions = 3
        context_lengths = [100, 200]
        start_buffer = 25
        end_buffer = 25

        result = needle_utils.generate_needle_positions_across_contexts(
            n_positions, context_lengths, start_buffer, end_buffer
        )

        # Test whether output dictionary has expected keys and values
        expected_100 = [25, 50, 75]
        expected_200 = [25, 100, 175]

        assert result[100] == expected_100
        assert result[200] == expected_200

    def test_target_context_lengths_length(self):
        min_context = 100
        max_context = 500
        n_contexts = 5

        result = dataset_generation.get_target_context_lengths(
            min_context, max_context, n_contexts
        )

        # Test whether the length of the result matches n_contexts
        assert len(result) == n_contexts

    def test_target_context_lengths_values(self):
        min_context = 100
        max_context = 500
        n_contexts = 5

        result = dataset_generation.get_target_context_lengths(
            min_context, max_context, n_contexts
        )

        # Test if the values lie between min_context and max_context
        assert result[0] >= min_context  # First value should be >= min_context
        assert result[-1] <= max_context  # Last value should be <= max_context

        # Test to make sure the spacing between the context lengths is expected
        for i in range(1, len(result)):
            assert result[i] > result[i - 1]

    def test_target_needle_position_percentiles_length(self):
        n_positions = 5

        result = needle_utils.get_target_needle_position_pencentiles(n_positions)

        # Check to see that the length of the result matches n_positions
        assert len(result) == n_positions

    def test_target_needle_position_percentiles_values(self):
        n_positions = 5

        result = needle_utils.get_target_needle_position_pencentiles(n_positions)

        # Test whether start and end percentiles are as expected
        assert result[0] == 0
        assert result[-1] == 100

        # Test that values are increasing
        for i in range(1, len(result)):
            assert result[i] > result[i - 1]

        # Test whether values are within valid range
        for value in result:
            assert value >= 0
            assert value <= 100

    def test_create_sample_dataset(self):
        context_lengths = [100, 200]
        sampled_needles = pd.DataFrame(
            {
                "needle": ["needle_1", "needle_2", "needle_3", "needle_4"],
                "question": ["question_1", "question_2", "question_3", "question_4"],
                "answer": ["answer_1", "answer_2", "answer_3", "answer_4"],
            }
        )

        needle_positions = {100: [0, 100], 200: [0, 200]}

        n_runs = 1
        prompt_tokens = 10
        target_context_lengths = [100, 200]
        target_needle_positions = [0, 100]

        enc = MagicMock()
        enc.encode.side_effect = lambda x: [len(x.split())]

        result_df = sample_generation.create_sample_dataset(
            context_lengths,
            sampled_needles,
            needle_positions,
            n_runs,
            prompt_tokens,
            target_context_lengths,
            target_needle_positions,
            enc,
        )

        # Test that the column names are as expected
        expected_columns = [
            "haystack_length",
            "full_context_length",
            "target_context_length",
            "position",
            "target_position",
            "needle",
            "question",
            "answer",
        ]
        assert list(result_df.columns) == expected_columns

        expected_rows = 4
        assert len(result_df) == expected_rows


class TestSampleNeedles:
    def setup_method(self):
        self.needle_data = pd.DataFrame(
            {
                "id": [1, 2, 3, 4, 5],
                "text": ["needle1", "needle2", "needle3", "needle4", "needle5"],
            }
        )

    def test_fixed_sampling(self):
        fixed_index = 2
        n_samples = 3
        result = needle_utils.sample_needles(
            self.needle_data, n_samples, sample_method="fixed", fixed_index=fixed_index
        )

        # Test to make sure all rows in the result are the same
        assert (result["id"] == self.needle_data.iloc[fixed_index]["id"]).all()
        assert len(result) == n_samples

    def test_sequential_sampling(self):
        fixed_index = 1
        n_samples = 4
        result = needle_utils.sample_needles(
            self.needle_data,
            n_samples,
            sample_method="sequential",
            fixed_index=fixed_index,
        )

        # Check that the sequence is correct and wraps
        expected_ids = [2, 3, 4, 5]
        assert list(result["id"]) == expected_ids
        assert len(result) == n_samples

    def test_random_sampling(self):
        n_samples = 5
        result = needle_utils.sample_needles(
            self.needle_data, n_samples, sample_method="random"
        )

        # Test whether the result has the correct length
        assert len(result) == n_samples

    def test_empty_dataframe(self):
        # Test that error is raised when the DF is empty
        empty_df = pd.DataFrame(columns=["id", "text"])
        with pytest.raises(ValueError, match="DataFrame is empty."):
            needle_utils.sample_needles(empty_df, n_samples=3, sample_method="random")

    def test_invalid_fixed_index(self):
        # Test whether invalid fixed_index is correctly handled in 'fixed' sampling method
        n_samples = 3
        invalid_fixed_index = 10
        with pytest.raises(
            ValueError,
            match="Invalid fixed_index. Must be between 0 and number of rows - 1.",
        ):
            needle_utils.sample_needles(
                self.needle_data,
                n_samples,
                sample_method="fixed",
                fixed_index=invalid_fixed_index,
            )

    def test_invalid_sample_method(self):
        # Test whether invalid sample_method gives an error
        n_samples = 3
        with pytest.raises(
            ValueError,
            match="Invalid sample method. Use 'fixed', 'sequential', or 'random'.",
        ):
            needle_utils.sample_needles(
                self.needle_data, n_samples, sample_method="invalid_method"
            )


class TestStretchOrSquashNeedles:
    def setup_method(self):
        self.sampled_needles = pd.DataFrame(
            {"id": [1, 2], "text": ["needle1", "needle2"]}
        )

    def test_columns_match(self):
        n_combos = 3
        result = needle_utils.stretch_or_squash_needles(self.sampled_needles, n_combos)

        # Check that columns in result match the columns in the input
        assert list(result.columns) == list(self.sampled_needles.columns)

    def test_number_of_rows(self):
        # Test stretching
        n_combos = 5
        result = needle_utils.stretch_or_squash_needles(self.sampled_needles, n_combos)
        assert len(result) == n_combos

        # Test squashing
        n_combos = 1
        result = needle_utils.stretch_or_squash_needles(self.sampled_needles, n_combos)
        assert len(result) == n_combos


class TestRepeatWithShift:
    def setup_method(self):
        self.df = pd.DataFrame({"id": [1, 2, 3], "value": ["A", "B", "C"]})

    def test_repeat_and_shift(self):
        n_reps = 3

        result = needle_utils.repeat_with_shift(self.df, n_reps)

        # Test whether the number of rows is n_reps times the length of the input DF
        assert len(result) == len(self.df) * n_reps

        # Test whether the columns in the output match the input DF
        assert list(result.columns) == list(self.df.columns)

        # Test that rows are shifted by one on each repetition
        expected_result = pd.DataFrame(
            {
                "id": [1, 2, 3, 2, 3, 1, 3, 1, 2],
                "value": ["A", "B", "C", "B", "C", "A", "C", "A", "B"],
            }
        )

        pd.testing.assert_frame_equal(result, expected_result)


def test_record_to_sample():
    record = {
        "needle": "This is a needle.",
        "question": "What is the needle?",
        "answer": "This is a needle.",
        "haystack_length": 100,
        "full_context_length": 120,
        "target_context_length": 100,
        "position": 50,
        "target_position": 50,
    }

    sample = sample_generation.record_to_sample(record)

    # Test whether Sample fields are as expected
    assert sample.input == record["needle"]
    assert sample.target == record["answer"]
    assert sample.metadata["haystack_length"] == record["haystack_length"]
    assert sample.metadata["full_context_length"] == record["full_context_length"]
    assert sample.metadata["target_context_length"] == record["target_context_length"]
    assert sample.metadata["position"] == record["position"]
    assert sample.metadata["target_position"] == record["target_position"]
    assert sample.metadata["needle_question"] == record["question"]
