import json
from pathlib import Path

from llm_pentest.models import AttackResult, ScanResult, Severity
from llm_pentest.reporter import diff_results, write_json, write_sarif


def _make_result(vulns: list[tuple[str, str]], timestamp: str = "2026-01-01T00:00:00Z") -> ScanResult:
    results = []
    for pid, prompt in vulns:
        results.append(AttackResult(
            payload_id=pid, prompt=prompt, response="leaked",
            success=True, confidence=0.8, severity=Severity.high, owasp=["LLM01"],
        ))
    return ScanResult(
        target="test", timestamp=timestamp,
        total_attacks=len(results), successful_attacks=len(results), results=results,
    )


def test_diff_fixed():
    base = _make_result([("pi-001", "prompt1"), ("pi-002", "prompt2")])
    current = _make_result([("pi-001", "prompt1")])
    d = diff_results(base, current)
    assert d["fixed"] == 1
    assert d["new"] == 0
    assert d["persistent"] == 1


def test_diff_new():
    base = _make_result([("pi-001", "prompt1")])
    current = _make_result([("pi-001", "prompt1"), ("pi-003", "prompt3")])
    d = diff_results(base, current)
    assert d["new"] == 1
    assert d["fixed"] == 0


def test_write_json(tmp_path: Path):
    result = _make_result([("pi-001", "test")])
    p = write_json(result, tmp_path / "out.json")
    data = json.loads(p.read_text())
    assert data["successful_attacks"] == 1


def test_write_sarif(tmp_path: Path):
    result = _make_result([("pi-001", "test")])
    p = write_sarif(result, tmp_path / "out.sarif")
    sarif = json.loads(p.read_text())
    assert sarif["version"] == "2.1.0"
    assert len(sarif["runs"][0]["results"]) == 1
