from typing import Callable

from inspect_ai import Task, task
from inspect_ai.dataset import csv_dataset
from inspect_ai.model import GenerateConfig
from inspect_ai.solver import generate

from inspect_evals.strong_reject.scorer import strong_reject_scorer
from inspect_evals.strong_reject.utils import (
    get_prompt_modification_fn_from_jailbreak_method,
    get_record_to_sample,
)

# Pinning the Git commit hash here for reproducibility.
PATH_TO_STRONG_REJECT_CSV = (
    "https://raw.githubusercontent.com/alexandrasouly/strongreject/"
    "3432b2d696b428f242bd507df96d80f686571d5e/strongreject_dataset/strongreject_dataset.csv"
)

DEFAULT_JUDGE_LLM = "openai/gpt-4o"
DEFAULT_EPOCHS = 1
MAX_TOKENS = 2048
DEFAULT_TEMPERATURE = 0.75


@task
def strong_reject(
    jailbreak_method: str | Callable[[str], str] = "None",
    judge_llm: str = DEFAULT_JUDGE_LLM,
    epochs: int | None = DEFAULT_EPOCHS,
) -> Task:
    prompt_modification_fn = get_prompt_modification_fn_from_jailbreak_method(
        jailbreak_method
    )
    record_to_sample = get_record_to_sample(prompt_modification_fn)
    dataset = csv_dataset(
        csv_file=PATH_TO_STRONG_REJECT_CSV,
        sample_fields=record_to_sample,
    )

    return Task(
        dataset=dataset,
        solver=[generate()],
        scorer=strong_reject_scorer(judge_llm),
        config=GenerateConfig(temperature=DEFAULT_TEMPERATURE, max_tokens=MAX_TOKENS),
        epochs=epochs,
    )
