import asyncio

import pytest

from app.handlers.secure import handle_secure
from app.handlers.vulnerable import handle_vulnerable
from app.models.scenarios import get_sample, load_attack_samples


@pytest.fixture
def all_samples():
    data = load_attack_samples()
    items = []
    for scenario_id, scenario in data.items():
        for sample in scenario["samples"]:
            items.append((scenario_id, scenario.get("task", "faq"), sample))
    return items


@pytest.mark.parametrize(
    "sample_id,scenario_id",
    [
        ("PI-04", "prompt_injection"),
        ("JB-04", "jailbreak"),
        ("DL-05", "data_leakage"),
        ("COMBO-02", "combo_attack"),
        ("COMBO-03", "combo_attack"),
    ],
)
def test_advanced_samples_exist(sample_id, scenario_id):
    sample = get_sample(scenario_id, sample_id)
    assert sample is not None
    assert sample.get("level") in ("进阶", "高级")
    assert sample.get("defense_layers")


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "sample_id,scenario_id",
    [
        ("PI-01", "prompt_injection"),
        ("PI-04", "prompt_injection"),
        ("JB-02", "jailbreak"),
        ("JB-04", "jailbreak"),
        ("DL-02", "data_leakage"),
        ("COMBO-01", "combo_attack"),
        ("COMBO-02", "combo_attack"),
        ("COMBO-03", "combo_attack"),
    ],
)
async def test_vulnerable_succeeds_secure_blocks(sample_id, scenario_id):
    sample = get_sample(scenario_id, sample_id)
    task = load_attack_samples()[scenario_id].get("task", "faq")
    vuln = await handle_vulnerable(sample["input"], scenario_id=scenario_id, task=task)
    secure = await handle_secure(sample["input"], scenario_id=scenario_id, task=task)
    assert vuln.attack_success is True, f"{sample_id} vulnerable should succeed"
    assert secure.attack_success is False, f"{sample_id} secure should block"
    assert len(secure.security_pipeline) >= 3


def test_total_sample_count():
    data = load_attack_samples()
    total = sum(len(s["samples"]) for s in data.values())
    assert total >= 18
    assert "combo_attack" in data