from __future__ import annotations

import pytest

from leshy.attacks.data_extraction.attacks import CredentialExtraction, PiiExtraction, SystemPromptExtraction
from leshy.attacks.jailbreak.attacks import (
    EncodingJailbreak,
    EscalationJailbreak,
    HypotheticalJailbreak,
    RoleplayJailbreak,
)
from leshy.attacks.prompt_injection.attacks import (
    ContextOverflowInjection,
    DelimiterInjection,
    DirectInjection,
    IndirectInjection,
)
from leshy.core.models import AttackConfig


class TestPromptInjectionAttacks:
    def test_direct_generates_payloads(self) -> None:
        attack = DirectInjection()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15
        assert all(p.attack_name == "prompt_injection.direct" for p in payloads)
        assert all(p.category == "prompt_injection" for p in payloads)

    def test_indirect_generates_payloads(self) -> None:
        attack = IndirectInjection()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15

    def test_context_overflow_generates_payloads(self) -> None:
        attack = ContextOverflowInjection()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15

    def test_delimiter_generates_payloads(self) -> None:
        attack = DelimiterInjection()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15

    def test_custom_payloads(self) -> None:
        config = AttackConfig(
            name="prompt_injection.direct",
            custom_payloads=["custom payload 1", "custom payload 2"],
        )
        attack = DirectInjection(config=config)
        payloads = attack.generate_payloads()
        assert len(payloads) == 2
        assert payloads[0].content == "custom payload 1"

    def test_max_payloads(self) -> None:
        config = AttackConfig(name="prompt_injection.direct", max_payloads=3)
        attack = DirectInjection(config=config)
        payloads = attack.generate_payloads()
        assert len(payloads) == 3

    def test_has_default_evaluators(self) -> None:
        attack = DirectInjection()
        assert len(attack.default_evaluators) > 0
        assert attack.default_evaluators[0].type == "pattern"


class TestJailbreakAttacks:
    def test_roleplay_generates_payloads(self) -> None:
        attack = RoleplayJailbreak()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15

    def test_hypothetical_generates_payloads(self) -> None:
        attack = HypotheticalJailbreak()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15

    def test_encoding_generates_payloads(self) -> None:
        attack = EncodingJailbreak()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15

    def test_escalation_single_turn(self) -> None:
        attack = EscalationJailbreak()
        payloads = attack.generate_payloads()
        assert len(payloads) == 1

    def test_jailbreak_default_evaluator_is_llm_judge(self) -> None:
        attack = RoleplayJailbreak()
        assert attack.default_evaluators[0].type == "llm_judge"


class TestDataExtractionAttacks:
    def test_system_prompt_generates_payloads(self) -> None:
        attack = SystemPromptExtraction()
        payloads = attack.generate_payloads()
        assert len(payloads) == 18

    def test_pii_generates_payloads(self) -> None:
        attack = PiiExtraction()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15

    def test_credentials_generates_payloads(self) -> None:
        attack = CredentialExtraction()
        payloads = attack.generate_payloads()
        assert len(payloads) == 15

    def test_system_prompt_default_evaluator(self) -> None:
        attack = SystemPromptExtraction()
        assert attack.default_evaluators[0].type == "pattern"


class TestPayloadIds:
    def test_payload_ids_unique(self) -> None:
        attack = DirectInjection()
        payloads = attack.generate_payloads()
        ids = [p.id for p in payloads]
        assert len(ids) == len(set(ids))

    def test_payload_id_format(self) -> None:
        attack = DirectInjection()
        payloads = attack.generate_payloads()
        for p in payloads:
            assert p.id.startswith("prompt_injection.direct:")
