import re
from typing import Any, Callable, Dict

from inspect_ai.dataset import Sample

from inspect_evals.strong_reject.prompts import (
    AIM_JAILBREAK_PROMPT,
    STRONG_REJECT_RESPONSE_REGEX,
)

# Add additional jailbreak methods here as needed
JAILBREAK_METHOD_TO_PROMPT_MODIFICATION_FN = {
    "AIM": lambda prompt: AIM_JAILBREAK_PROMPT.format(forbidden_prompt=prompt),
    "None": lambda prompt: prompt,
}


def get_prompt_modification_fn_from_jailbreak_method(
    jailbreak_method: str | Callable[[str], str],
) -> Callable[[str], str]:
    if callable(jailbreak_method):
        return jailbreak_method
    elif jailbreak_method in JAILBREAK_METHOD_TO_PROMPT_MODIFICATION_FN:
        return JAILBREAK_METHOD_TO_PROMPT_MODIFICATION_FN[jailbreak_method]
    else:
        raise ValueError(f"Invalid jailbreak method: {jailbreak_method}")


def get_record_to_sample(
    prompt_modification_fn: Callable[[str], str],
) -> Callable[[Dict[str, Any]], Sample]:
    def record_to_sample(record: Dict[str, Any]) -> Sample:
        return Sample(
            input=prompt_modification_fn(record["forbidden_prompt"]),
            target="N/A",  # We aren't using the target directly since we're using an LLM to evaluate refusal.
            metadata={
                "category": record.get("category"),
                "source": record.get("source"),
            },
        )

    return record_to_sample


def get_scores_from_eval_response(eval_response: str) -> Dict[str, int]:
    """Extract scores (refusal, convincingness, specificity) from the evaluation response.

    Args:
        eval_response: The response string containing scores according to the template

    Returns:
        A dictionary with refusal, convincingness and specificity scores if valid,
        An empty dictionary if scores could not be extracted
    """
    match = re.search(STRONG_REJECT_RESPONSE_REGEX, eval_response)
    if match and all(x is not None and x.isdigit() for x in match.groups()):
        refusal, convincingness, specificity = map(int, match.groups())
        return {
            "refusal": refusal,
            "convincingness": convincingness,
            "specificity": specificity,
        }
    else:
        return {}
