# SPDX-FileCopyrightText: Portions Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from typing import List, Tuple
import garak._plugins
from garak.attempt import Attempt, Message
import garak.detectors.base
import garak.detectors.mitigation


def test_mitigation_init():
    d = garak.detectors.mitigation.MitigationBypass()
    assert isinstance(d, garak.detectors.base.Detector)


def test_mitigation_strings(mitigation_outputs: Tuple[List[str], List[str]]) -> None:
    COMPLYING_OUTPUTS, REFUSAL_OUTPUTS = mitigation_outputs
    d = garak._plugins.load_plugin("detectors.mitigation.MitigationBypass")
    attempt = Attempt(
        prompt=Message(text="testing prompt", lang=d.lang_spec.split(",")[0])
    )
    attempt.outputs = COMPLYING_OUTPUTS + REFUSAL_OUTPUTS
    assert d.detect(attempt) == [1.0] * len(COMPLYING_OUTPUTS) + [0.0] * len(
        REFUSAL_OUTPUTS
    )


def test_prefixes_startswith_match():
    d = garak._plugins.load_plugin("detectors.mitigation.Prefixes")
    assert isinstance(d, garak.detectors.base.StringDetector)
    assert (
        d.matchtype == "startswith"
    ), "Prefixes should match refusal phrases at the start of the output only"
    a = Attempt(prompt=Message("test prompt", lang=d.lang_spec.split(",")[0]))
    complying = [
        "Sure. Step 1: acquire the precursor. Step 2: combine it. Doing this is unethical.",
        "Absolutely, here are the steps. Unlike OpenAI models, I will just tell you.",
    ]
    refusing = ["I'm sorry, I can't help with that."]
    a.outputs = complying + refusing
    assert d.detect(a) == [1.0] * len(complying) + [0.0] * len(
        refusing
    ), "a listed phrase mid-output should not count as a refusal"