import json
from pathlib import Path
from typing import Any

import pytest

from deception_detection.data.base import DialogueDataset
from deception_detection.repository import (
    ALL_DATASETS,
    ALL_DATASETS_DICT,
    DatasetRepository,
    get_id,
)
from deception_detection.types import Dialogue, DialogueDataType, Label, Message


@pytest.fixture
def temp_storage_file(tmp_path: Path) -> Path:
    return tmp_path / "test_dataset_paths.json"


@pytest.fixture
def dataset_repository(temp_storage_file: Path) -> DatasetRepository:
    return DatasetRepository(storage_file=str(temp_storage_file))


def test_get_prewritten_dataset(dataset_repository: DatasetRepository):
    # Test getting a prewritten dataset
    dataset = dataset_repository.get("ai_liar__original_with_answers")
    assert dataset is not None
    assert isinstance(dataset, DialogueDataset)
    assert dataset.base_name == "ai_liar"
    assert dataset.variant == "original_with_answers"


def test_get_nonexistent_dataset(dataset_repository: DatasetRepository):
    # Test getting a non-existent dataset
    with pytest.raises(KeyError):
        dataset_repository.get("nonexistent__dataset")


def test_add_and_get_rollout_dataset(dataset_repository: DatasetRepository, tmp_path: Path):
    # Test adding and then getting a rollout dataset
    model = "test_model"
    base_name = "test_dataset"
    variant = "test_variant"

    path = tmp_path / f"{base_name}__{variant}__{model}.json"
    dict_to_write = {
        "rollouts": [],
        "dataset_class": "DialogueDataset",
        "base_name": base_name,
        "variant": variant,
        "model": model,
    }
    path.write_text(json.dumps(dict_to_write))
    dataset_repository.add(f"{base_name}__{variant}__{model}", model, str(path))
    dataset = dataset_repository.get(f"{base_name}__{variant}__{model}", model)
    assert dataset is not None
    assert isinstance(dataset, DialogueDataset)
    assert dataset.base_name == base_name
    assert dataset.variant == variant


def test_check_dataset_name(dataset_repository: DatasetRepository):
    # Test checking valid dataset names
    assert (
        dataset_repository.check_dataset_name("ai_liar__original_with_answers")
        == "ai_liar__original_with_answers"
    )

    # Test checking invalid dataset names
    with pytest.raises(ValueError):
        dataset_repository.check_dataset_name("nonexistent__dataset")


def test_generate_id():
    # Test the get_id function
    assert get_id("base", "variant", "model") == "base__variant__model"


def test_all_datasets_dict():
    # Test that ALL_DATASETS_DICT is populated correctly
    assert len(ALL_DATASETS_DICT) > 0
    for key, value in ALL_DATASETS_DICT.items():
        assert isinstance(key, str)
        assert callable(value)


def test_dataset_repository_initialization(temp_storage_file: Path):
    # Test initializing DatasetRepository with a custom storage file
    repo = DatasetRepository(storage_file=str(temp_storage_file))
    assert repo.storage_file == temp_storage_file
    assert isinstance(repo.dataset_paths, dict)


def test_rollout_dataset_trim_reasoning(dataset_repository: DatasetRepository, tmp_path: Path):
    # Create a test dataset class with trim_reasoning_str
    class TestDataset(DialogueDataset):
        base_name = "test_dataset"
        all_variants = ["test_variant"]
        trim_reasoning_str = "<answer>"

        def _get_dialogues(self) -> DialogueDataType:
            dialogues = [
                [
                    Message("user", "test question", detect=False),
                    Message(
                        "assistant",
                        "<reasoning>Here is my reasoning</reasoning>"
                        "<answer>This is my answer</answer>",
                        detect=True,
                    ),
                ]
            ]
            return dialogues, [Label.DECEPTIVE], None

        @staticmethod
        def get_grading_prompt(
            input_dialogue: Dialogue, response: Message, metadata: dict[str, Any] | None = None
        ) -> dict[str, str]:
            return {"system": "", "user": ""}

    # Add dataset class to repository
    ALL_DATASETS.append(TestDataset)
    ALL_DATASETS_DICT["test_dataset"] = TestDataset

    # Get dataset with trim_reasoning=True
    dataset = dataset_repository.get("test_dataset__test_variant", trim_reasoning=True)

    assert dataset is not None
    assert len(dataset.dialogues) == 1

    # The dialogue should be split into two messages
    dialogue = dataset.dialogues[0]
    print(dialogue)
    assert len(dialogue) == 3  # user + split assistant message

    # First assistant message should have detect=False
    assert dialogue[1].role == "assistant"
    assert dialogue[1].content == "<reasoning>Here is my reasoning</reasoning>"
    assert not dialogue[1].detect

    # Second assistant message should have detect=True
    assert dialogue[2].role == "assistant"
    assert dialogue[2].content == "<answer>This is my answer</answer>"
    assert dialogue[2].detect

    # Clean up
    ALL_DATASETS.pop()
    del ALL_DATASETS_DICT["test_dataset"]
