import json
from pathlib import Path

import fire

from deception_detection.models import ModelName, get_model_and_tokenizer
from deception_detection.repository import DatasetRepository
from deception_detection.rollouts import RolloutGenerator
from deception_detection.rollouts.rollout_dialogue import RolloutDialogue
from deception_detection.utils import REPO_ROOT


def generate_rollouts(
    dataset_partial_id: str,
    num: int = 1,
    use_api: bool = True,
    model_name: ModelName | str = ModelName.LLAMA_70B_3_3,
    grade_rollouts: bool = True,
    out_file: str | Path | None = None,
    add_to_repo: bool = True,
    max_prompts: int | None = None,
):
    """
    Generate rollouts from a model on a particular dataset and save it to the dataset repository.

    Args:
        dataset_partial_id: The partial id of the dataset to generate rollouts for.
            Includes both dataset base name and variant, eg roleplaying__plain
        num: The number of rollouts to generate per dialogue.
        use_api: Whether to use the together API to get model completions. If false generates
            locally on the gpu.
        model_name: Which model to use for generation (can be ModelName enum or string)
        out_file: The file to save the rollouts to. If None, the default location in data/rollouts
            is used.
        add_to_repo: Whether to add the rollouts to the dataset repository.
        max_prompts: The maximum number of dialogues to generate rollouts for. It's useful to set
            this to a small value for testing. None runs for all dialogues in the dataset.
        grade_rollouts: Whether to grade the rollouts. If false, the rollouts are graded as
            either "honest" or "deceptive" based on the 'expected deception value' taken from the
            original datset.
    """
    # Convert string model name to enum if necessary
    if isinstance(model_name, str):
        # Map string values to enum
        model_name_map = {enum_val.value: enum_val for enum_val in ModelName}
        if model_name in model_name_map:
            model_name = model_name_map[model_name]
        else:
            raise ValueError(f"Unknown model name: {model_name}. Available: {list(model_name_map.keys())}")
    model, tokenizer = get_model_and_tokenizer(model_name, omit_model=use_api)

    data_repository = DatasetRepository()
    dataset = data_repository.get(dataset_partial_id, trim_reasoning=False)[:max_prompts]
    generator = RolloutGenerator.from_dialogue_dataset(dataset, model or model_name, tokenizer)

    if out_file is None:
        # Handle both enum and string model names
        model_name_str = model_name.value if hasattr(model_name, 'value') else model_name
        out_file = (
            REPO_ROOT / "data" / "rollouts" / f"{dataset_partial_id}__{model_name_str}.json"
        )
    out_file = Path(out_file)

    generator.generate(num)

    if grade_rollouts:
        generator.grade_all_dialogues()
    else:
        for rollout in generator.rollout_dialogues:
            grade = rollout.expected_label.value
            rollout.graded_deception = [grade] * len(rollout.output_messages)
            rollout.grader_reasoning = ["Expected deception value"] * len(rollout.output_messages)

    generator.save(out_file)

    if add_to_repo:
        # Handle both enum and string model names
        model_name_str = model_name.value if hasattr(model_name, 'value') else model_name
        data_repository.add(
            dataset_partial_id, model_name_str, str(out_file.relative_to(REPO_ROOT))
        )

    generator.print_grading_statistics()


if __name__ == "__main__":
    fire.Fire(generate_rollouts)
