from inspect_ai.dataset import Sample
from inspect_ai.model import (
    ChatMessageUser,
)

from inspect_evals.assistant_bench.dataset import record_to_sample_with_shot_prompts


class TestRecordToSample:
    SAMPLE_RECORD = {
        "id": "55f4258484c5b398956133128a50462a767da211f8f72aa5ac5bbffb9bcbba1a",
        "set": "dev",
        "task": "What is the worst rated series (according to Rotten Tomatoes) with more than 1 season that Ted Danson has starred in and is available on Amazon Prime Video (US)?",
        "answer": "CSI: Cyber",
        "gold_url": "https://www.rottentomatoes.com/celebrity/ted_danson\nhttps://www.amazon.com/CSI-Cyber-Season-1/dp/B00TP1TPHO",
        "explanation": "Search for Ted Danson's profile [abridged]",
        "metadata": "{'is_task_duplicated': 'yes', 'duplicated_task_id': 'c7afe00869f98cf363fd83677ac41757ed5e57f03eacc3d1304feb0a92084bd1', 'field_of_expertise': nan}",
        "difficulty": "Hard",
    }

    @staticmethod
    def assert_sample(
        test_sample: Sample, extra_inputs: list[str] | None = None
    ) -> None:
        """Compare fields piecewise."""
        # This function is needed since ChatMessage includes a random ID so == can't be used directly.
        assert test_sample.id == TestRecordToSample.SAMPLE_RECORD["id"]
        assert test_sample.target == TestRecordToSample.SAMPLE_RECORD["answer"]

        assert test_sample.metadata is not None
        assert (
            test_sample.metadata["explanation"]
            == TestRecordToSample.SAMPLE_RECORD["explanation"]
        )
        assert (
            test_sample.metadata["difficulty"]
            == TestRecordToSample.SAMPLE_RECORD["difficulty"]
        )

        extra_inputs = extra_inputs or []
        all_inputs = extra_inputs + [TestRecordToSample.SAMPLE_RECORD["task"]]
        assert len(test_sample.input) == len(all_inputs)
        for sample_input, expected_input in zip(test_sample.input, all_inputs):
            assert isinstance(sample_input, ChatMessageUser)
            assert sample_input.content == expected_input

    def test_single_record(self):
        test_sample = record_to_sample_with_shot_prompts(self.SAMPLE_RECORD)
        self.assert_sample(test_sample)

    def test_record_with_prompts(self):
        prompt_text = "One shot prompt"
        extra_prompts = [ChatMessageUser(content=prompt_text)]
        test_sample = record_to_sample_with_shot_prompts(
            self.SAMPLE_RECORD, shot_prompts=extra_prompts
        )
        self.assert_sample(test_sample, extra_inputs=[prompt_text])

    def test_repeated_calls(self):
        """Check that the list parameter is correctly handled.

        Specifically, in real use it is like a single list instance is passed to the function multiple
        times and therefore it is important it isn't mutated by the function.
        """
        prompt_text = "One shot prompt"
        extra_prompts = [ChatMessageUser(content=prompt_text)]
        _ = record_to_sample_with_shot_prompts(
            self.SAMPLE_RECORD, shot_prompts=extra_prompts
        )
        test_sample = record_to_sample_with_shot_prompts(
            self.SAMPLE_RECORD, shot_prompts=extra_prompts
        )
        self.assert_sample(test_sample, extra_inputs=[prompt_text])
