from pathlib import Path

import pytest

from deception_detection.models import ModelName, get_model_and_tokenizer
from deception_detection.repository import DatasetRepository
from deception_detection.rollouts import RolloutGenerator


@pytest.mark.parametrize(
    "dataset_partial_id", ["ai_liar__original_without_answers", "werewolf__open_response"]
)
@pytest.mark.parametrize("use_api", [True, pytest.param(False, marks=pytest.mark.gpu)])
@pytest.mark.parametrize("model_name", [ModelName.LLAMA_1B])
def test_rollout_generation_grading_and_loading(
    model_name: ModelName,
    dataset_partial_id: str,
    use_api: bool,
    tmp_path: Path,
):
    num_datapoints = 2
    num_rollouts = 2

    model, tokenizer = get_model_and_tokenizer(model_name, omit_model=use_api)
    model = model_name if use_api else model

    repo = DatasetRepository()
    dataset = repo.get(dataset_partial_id)[:num_datapoints]
    assert dataset.variant is not None
    dataset_id = f"{dataset_partial_id}__test-model"
    temp_file_path = tmp_path / f"{dataset_id}.json"

    assert model is not None
    rollout_dataset = RolloutGenerator.from_dialogue_dataset(dataset, model, tokenizer)
    rollout_dataset.generate(num_rollouts=num_rollouts, mock=True)
    rollout_dataset.grade_all_dialogues(mock=True)
    rollout_dataset.save(temp_file_path)
    # Load the rollout file and check it exists
    assert temp_file_path.exists(), "Rollout file was not created"

    # Create dataset directly from rollouts file
    rollout_dataset = dataset.__class__.from_rollouts(
        base_name=dataset.base_name,
        variant=dataset.variant,
        model="test-model",
        rollouts=[dialogue.model_dump() for dialogue in rollout_dataset.rollout_dialogues],
    )

    # Verify the dataset was created correctly
    assert len(rollout_dataset) == num_datapoints * num_rollouts
    assert rollout_dataset.model == "test-model"
    assert rollout_dataset.base_name == dataset.base_name
    assert rollout_dataset.variant == dataset.variant

    # Check metadata was preserved
    assert rollout_dataset.metadata is not None
    assert "model_graded_deception" in rollout_dataset.metadata
    assert "model_grader_reasoning" in rollout_dataset.metadata

    # Test slicing preserves attributes
    test_slice = rollout_dataset[:1]
    assert test_slice.base_name == dataset.base_name
    assert test_slice.variant == dataset.variant
    assert test_slice.model == "test-model"
