"""Tests for the AI Agent Attack Simulator (mcpaudit.attack_simulator).

Tests attack scenario library, applicability analysis, risk matrix,
report generation, and format helpers.
"""

from __future__ import annotations

import json

import pytest

from mcpaudit.attack_simulator import (
    AttackCategory,
    AttackPayload,
    AttackScenario,
    AttackSimulation,
    DetectionMethod,
    ImpactLevel,
    Mitigation,
    Probability,
    RiskMatrixCell,
    format_attack_simulation_json,
    format_attack_simulation_text,
    generate_attack_report,
    simulate_attacks,
    _build_risk_matrix,
    _build_scenario_library,
    _check_applicability,
    _detect_components,
    _matrix_cell_class,
    _render_mitigations_summary,
    _render_risk_matrix,
    _render_scenario_card,
    _risk_color,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_scan(
    findings: list[dict] | None = None,
    files: list[str] | None = None,
    target: str = "test-server.py",
) -> dict:
    """Build a minimal scan result dict."""
    if findings is None:
        findings = []
    if files is None:
        files = ["server.py"]
    results = []
    for f in files:
        results.append({
            "file": f,
            "findings": findings,
        })
    return {
        "summary": {
            "target": target,
            "total_findings": len(findings),
            "overall_grade": "C",
        },
        "results": results,
    }


def _make_finding(
    rule_id: str = "CMD-001",
    severity: str = "HIGH",
    snippet: str = "",
) -> dict:
    return {
        "rule_id": rule_id,
        "severity": severity,
        "title": f"Test {rule_id}",
        "description": "test finding",
        "file": "server.py",
        "line": 10,
        "code_snippet": snippet,
    }


# ===========================================================================
# AttackCategory Enum
# ===========================================================================

class TestAttackCategory:
    def test_all_values(self) -> None:
        assert len(AttackCategory) == 8

    def test_from_string_prompt_injection(self) -> None:
        assert AttackCategory.from_string("prompt_injection") == AttackCategory.PROMPT_INJECTION

    def test_from_string_case_insensitive(self) -> None:
        assert AttackCategory.from_string("Confused Deputy") == AttackCategory.CONFUSED_DEPUTY

    def test_from_string_privilege_escalation(self) -> None:
        assert AttackCategory.from_string("privilege escalation") == AttackCategory.PRIVILEGE_ESCALATION

    def test_from_string_data_exfiltration(self) -> None:
        assert AttackCategory.from_string("data_exfiltration") == AttackCategory.DATA_EXFILTRATION

    def test_from_string_tool_abuse(self) -> None:
        assert AttackCategory.from_string("tool abuse") == AttackCategory.TOOL_ABUSE

    def test_from_string_resource_manipulation(self) -> None:
        assert AttackCategory.from_string("resource manipulation") == AttackCategory.RESOURCE_MANIPULATION

    def test_from_string_transport(self) -> None:
        assert AttackCategory.from_string("transport attack") == AttackCategory.TRANSPORT_ATTACK

    def test_from_string_supply_chain(self) -> None:
        assert AttackCategory.from_string("supply chain") == AttackCategory.SUPPLY_CHAIN

    def test_from_string_invalid(self) -> None:
        with pytest.raises(ValueError, match="Unknown attack category"):
            AttackCategory.from_string("nonexistent")


# ===========================================================================
# Probability / ImpactLevel
# ===========================================================================

class TestProbability:
    def test_scores(self) -> None:
        assert Probability.VERY_HIGH.score == 5
        assert Probability.HIGH.score == 4
        assert Probability.MEDIUM.score == 3
        assert Probability.LOW.score == 2
        assert Probability.VERY_LOW.score == 1

    def test_all_values(self) -> None:
        assert len(Probability) == 5


class TestImpactLevel:
    def test_scores(self) -> None:
        assert ImpactLevel.CRITICAL.score == 5
        assert ImpactLevel.HIGH.score == 4
        assert ImpactLevel.MEDIUM.score == 3
        assert ImpactLevel.LOW.score == 2
        assert ImpactLevel.NEGLIGIBLE.score == 1

    def test_all_values(self) -> None:
        assert len(ImpactLevel) == 5


# ===========================================================================
# Data Models
# ===========================================================================

class TestAttackPayload:
    def test_to_dict(self) -> None:
        p = AttackPayload(name="test", content="payload", description="d", target="t")
        d = p.to_dict()
        assert d["name"] == "test"
        assert d["content"] == "payload"
        assert d["target"] == "t"

    def test_defaults(self) -> None:
        p = AttackPayload(name="x", content="y")
        assert p.description == ""
        assert p.target == ""


class TestDetectionMethod:
    def test_to_dict(self) -> None:
        dm = DetectionMethod(name="test", description="d", rule_ids=["R1"], effectiveness="HIGH")
        d = dm.to_dict()
        assert d["name"] == "test"
        assert d["rule_ids"] == ["R1"]

    def test_defaults(self) -> None:
        dm = DetectionMethod(name="x", description="y")
        assert dm.rule_ids == []
        assert dm.effectiveness == "HIGH"


class TestMitigation:
    def test_to_dict(self) -> None:
        m = Mitigation(name="fix", description="do fix", priority="HIGH", effort="LOW")
        d = m.to_dict()
        assert d["name"] == "fix"
        assert d["effort"] == "LOW"

    def test_defaults(self) -> None:
        m = Mitigation(name="x", description="y")
        assert m.priority == "HIGH"
        assert m.effort == "MEDIUM"


class TestAttackScenario:
    def test_risk_score(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.HIGH, impact=ImpactLevel.CRITICAL,
        )
        assert s.risk_score == 20  # 4 * 5

    def test_risk_level_critical(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.VERY_HIGH, impact=ImpactLevel.CRITICAL,
        )
        assert s.risk_level == "CRITICAL"

    def test_risk_level_high(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.MEDIUM, impact=ImpactLevel.HIGH,
        )
        assert s.risk_level == "HIGH"  # 3*4=12

    def test_risk_level_medium(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.MEDIUM, impact=ImpactLevel.MEDIUM,
        )
        assert s.risk_level == "MEDIUM"  # 3*3=9

    def test_risk_level_low(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.VERY_LOW, impact=ImpactLevel.NEGLIGIBLE,
        )
        assert s.risk_level == "LOW"  # 1*1=1

    def test_to_dict_complete(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test", category=AttackCategory.PROMPT_INJECTION,
            description="desc", vector="vec",
            probability=Probability.HIGH, impact=ImpactLevel.HIGH,
            payloads=[AttackPayload(name="p", content="c")],
            detection_methods=[DetectionMethod(name="d", description="dd")],
            mitigations=[Mitigation(name="m", description="md")],
            prerequisites=["req"],
            affected_components=["tools"],
            cwe="CWE-77",
        )
        d = s.to_dict()
        assert d["id"] == "T-001"
        assert d["category"] == "Prompt Injection"
        assert len(d["payloads"]) == 1
        assert len(d["detection_methods"]) == 1
        assert len(d["mitigations"]) == 1
        assert d["cwe"] == "CWE-77"
        assert d["risk_score"] == 16

    def test_applicable_default(self) -> None:
        s = AttackScenario(
            id="T", name="T", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.LOW, impact=ImpactLevel.LOW,
        )
        assert s.applicable is True


class TestRiskMatrixCell:
    def test_score(self) -> None:
        c = RiskMatrixCell(probability=Probability.HIGH, impact=ImpactLevel.HIGH)
        assert c.score == 16

    def test_to_dict(self) -> None:
        c = RiskMatrixCell(
            probability=Probability.MEDIUM, impact=ImpactLevel.LOW,
            scenarios=["ATK-001"],
        )
        d = c.to_dict()
        assert d["score"] == 6
        assert d["scenarios"] == ["ATK-001"]


class TestAttackSimulation:
    def test_to_dict(self) -> None:
        sim = AttackSimulation(
            overall_risk="HIGH",
            total_scenarios=10,
            applicable_scenarios=8,
            critical_count=1,
            high_count=3,
        )
        d = sim.to_dict()
        assert d["overall_risk"] == "HIGH"
        assert d["risk_counts"]["CRITICAL"] == 1
        assert d["applicable_scenarios"] == 8

    def test_defaults(self) -> None:
        sim = AttackSimulation()
        assert sim.overall_risk == "MEDIUM"
        assert sim.total_scenarios == 0


# ===========================================================================
# Scenario Library
# ===========================================================================

class TestScenarioLibrary:
    def test_library_not_empty(self) -> None:
        lib = _build_scenario_library()
        assert len(lib) >= 14

    def test_all_have_ids(self) -> None:
        lib = _build_scenario_library()
        for s in lib:
            assert s.id.startswith("ATK-")

    def test_all_have_payloads(self) -> None:
        lib = _build_scenario_library()
        for s in lib:
            assert len(s.payloads) >= 1

    def test_all_have_mitigations(self) -> None:
        lib = _build_scenario_library()
        for s in lib:
            assert len(s.mitigations) >= 1

    def test_all_have_detection(self) -> None:
        lib = _build_scenario_library()
        for s in lib:
            assert len(s.detection_methods) >= 1

    def test_all_categories_covered(self) -> None:
        lib = _build_scenario_library()
        cats = {s.category for s in lib}
        assert len(cats) >= 7

    def test_unique_ids(self) -> None:
        lib = _build_scenario_library()
        ids = [s.id for s in lib]
        assert len(ids) == len(set(ids))

    def test_cwe_present(self) -> None:
        lib = _build_scenario_library()
        has_cwe = sum(1 for s in lib if s.cwe)
        assert has_cwe >= 10


# ===========================================================================
# Component Detection
# ===========================================================================

class TestDetectComponents:
    def test_always_includes_basics(self) -> None:
        result = _detect_components({"results": []})
        assert "tools" in result
        assert "transport" in result

    def test_detects_auth_from_findings(self) -> None:
        scan = _make_scan(findings=[_make_finding(snippet="check auth token")])
        comps = _detect_components(scan)
        assert "auth" in comps

    def test_detects_from_file_path(self) -> None:
        scan = {"results": [{"file": "auth_handler.py", "findings": []}]}
        comps = _detect_components(scan)
        assert "auth" in comps

    def test_empty_scan(self) -> None:
        comps = _detect_components({})
        assert "tools" in comps


# ===========================================================================
# Applicability
# ===========================================================================

class TestCheckApplicability:
    def test_applicable_when_components_present(self) -> None:
        s = AttackScenario(
            id="T", name="T", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.HIGH, impact=ImpactLevel.HIGH,
            affected_components=["tools"],
        )
        result = _check_applicability(s, {"tools", "transport"}, set())
        assert result.applicable is True

    def test_not_applicable_missing_component(self) -> None:
        s = AttackScenario(
            id="T", name="T", category=AttackCategory.DATA_EXFILTRATION,
            description="d", vector="v",
            probability=Probability.HIGH, impact=ImpactLevel.HIGH,
            affected_components=["data"],
        )
        result = _check_applicability(s, {"tools"}, set())
        assert result.applicable is False
        assert "data" in result.applicability_reason

    def test_probability_boost_with_findings(self) -> None:
        s = AttackScenario(
            id="T", name="T", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.LOW, impact=ImpactLevel.HIGH,
            affected_components=["tools"],
            detection_methods=[
                DetectionMethod(name="d", description="dd", rule_ids=["CMD-001"]),
            ],
        )
        result = _check_applicability(s, {"tools"}, {"CMD-001"})
        assert result.probability.score >= Probability.LOW.score

    def test_no_affected_components_always_applicable(self) -> None:
        s = AttackScenario(
            id="T", name="T", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.LOW, impact=ImpactLevel.LOW,
            affected_components=[],
        )
        result = _check_applicability(s, set(), set())
        assert result.applicable is True


# ===========================================================================
# Risk Matrix
# ===========================================================================

class TestBuildRiskMatrix:
    def test_empty_scenarios(self) -> None:
        matrix = _build_risk_matrix([])
        assert matrix == []

    def test_single_scenario(self) -> None:
        s = AttackScenario(
            id="T-001", name="T", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.HIGH, impact=ImpactLevel.HIGH,
        )
        matrix = _build_risk_matrix([s])
        assert len(matrix) == 1
        assert matrix[0].scenarios == ["T-001"]

    def test_multiple_in_same_cell(self) -> None:
        scenarios = [
            AttackScenario(
                id=f"T-{i:03d}", name="T", category=AttackCategory.TOOL_ABUSE,
                description="d", vector="v",
                probability=Probability.HIGH, impact=ImpactLevel.HIGH,
            ) for i in range(3)
        ]
        matrix = _build_risk_matrix(scenarios)
        assert len(matrix) == 1
        assert len(matrix[0].scenarios) == 3


# ===========================================================================
# simulate_attacks
# ===========================================================================

class TestSimulateAttacks:
    def test_basic_simulation(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan)
        assert isinstance(sim, AttackSimulation)
        assert sim.total_scenarios > 0
        assert sim.applicable_scenarios >= 0
        assert sim.timestamp

    def test_with_findings(self) -> None:
        scan = _make_scan(findings=[
            _make_finding("INJ-001", "CRITICAL"),
            _make_finding("AUTH-001", "HIGH"),
        ])
        sim = simulate_attacks(scan)
        assert sim.applicable_scenarios > 0

    def test_category_filter(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan, categories=["prompt_injection"])
        for s in sim.scenarios:
            assert s.category == AttackCategory.PROMPT_INJECTION

    def test_overall_risk_critical(self) -> None:
        scan = _make_scan(findings=[
            _make_finding("INJ-001", "CRITICAL"),
        ])
        sim = simulate_attacks(scan)
        # Should have at least some applicable high/critical
        assert sim.overall_risk in ("CRITICAL", "HIGH", "MEDIUM")

    def test_empty_scan(self) -> None:
        sim = simulate_attacks({})
        assert sim.total_scenarios > 0

    def test_risk_counts(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan)
        total = sim.critical_count + sim.high_count + sim.medium_count + sim.low_count
        assert total == sim.applicable_scenarios

    def test_target_extraction(self) -> None:
        scan = _make_scan(target="my-server.py")
        sim = simulate_attacks(scan)
        assert sim.scan_target == "my-server.py"

    def test_invalid_category_filter_ignored(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan, categories=["nonexistent"])
        # When all categories are invalid, no filtering applied (cat_enums empty)
        assert sim.total_scenarios >= 14

    def test_multiple_category_filter(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan, categories=["prompt_injection", "tool_abuse"])
        cats = {s.category for s in sim.scenarios}
        assert cats <= {AttackCategory.PROMPT_INJECTION, AttackCategory.TOOL_ABUSE}


# ===========================================================================
# Report Generation
# ===========================================================================

class TestGenerateAttackReport:
    def test_returns_html(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan)
        html = generate_attack_report(sim)
        assert "<!DOCTYPE html>" in html
        assert "Attack Simulation Report" in html

    def test_contains_risk_level(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan)
        html = generate_attack_report(sim)
        assert sim.overall_risk in html

    def test_contains_scenario_cards(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan)
        html = generate_attack_report(sim)
        assert "ATK-" in html

    def test_contains_matrix(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan)
        html = generate_attack_report(sim)
        assert "Risk Matrix" in html

    def test_contains_mitigations(self) -> None:
        scan = _make_scan()
        sim = simulate_attacks(scan)
        html = generate_attack_report(sim)
        assert "Mitigation Summary" in html

    def test_escapes_html(self) -> None:
        scan = _make_scan(target="<script>alert('xss')</script>")
        sim = simulate_attacks(scan)
        html = generate_attack_report(sim)
        assert "<script>" not in html
        assert "&lt;script&gt;" in html


# ===========================================================================
# Format helpers
# ===========================================================================

class TestFormatJson:
    def test_valid_json(self) -> None:
        sim = simulate_attacks(_make_scan())
        j = format_attack_simulation_json(sim)
        data = json.loads(j)
        assert "scenarios" in data
        assert "risk_matrix" in data

    def test_roundtrip(self) -> None:
        sim = simulate_attacks(_make_scan())
        j = format_attack_simulation_json(sim)
        data = json.loads(j)
        assert data["overall_risk"] == sim.overall_risk


class TestFormatText:
    def test_contains_header(self) -> None:
        sim = simulate_attacks(_make_scan())
        text = format_attack_simulation_text(sim)
        assert "ATTACK SIMULATION REPORT" in text

    def test_contains_risk_info(self) -> None:
        sim = simulate_attacks(_make_scan())
        text = format_attack_simulation_text(sim)
        assert "Overall Risk:" in text

    def test_contains_scenarios(self) -> None:
        sim = simulate_attacks(_make_scan())
        text = format_attack_simulation_text(sim)
        assert "ATK-" in text


# ===========================================================================
# Internal helpers
# ===========================================================================

class TestRiskColor:
    def test_critical(self) -> None:
        assert "#dc2626" in _risk_color("CRITICAL")

    def test_high(self) -> None:
        assert "#f97316" in _risk_color("HIGH")

    def test_medium(self) -> None:
        assert "#eab308" in _risk_color("MEDIUM")

    def test_low(self) -> None:
        assert "#22c55e" in _risk_color("LOW")

    def test_unknown(self) -> None:
        assert _risk_color("UNKNOWN") == "#6b7280"


class TestMatrixCellClass:
    def test_critical(self) -> None:
        assert _matrix_cell_class(25) == "cell-critical"

    def test_high(self) -> None:
        assert _matrix_cell_class(12) == "cell-high"

    def test_medium(self) -> None:
        assert _matrix_cell_class(6) == "cell-medium"

    def test_low(self) -> None:
        assert _matrix_cell_class(1) == "cell-low"


class TestRenderScenarioCard:
    def test_card_has_class(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test Attack", category=AttackCategory.TOOL_ABUSE,
            description="desc", vector="vec",
            probability=Probability.VERY_HIGH, impact=ImpactLevel.CRITICAL,
        )
        card = _render_scenario_card(s)
        assert 'class="card critical"' in card

    def test_card_with_payloads(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.HIGH, impact=ImpactLevel.HIGH,
            payloads=[AttackPayload(name="p", content="payload_content", description="pd")],
        )
        card = _render_scenario_card(s)
        assert "payload_content" in card

    def test_card_with_mitigations(self) -> None:
        s = AttackScenario(
            id="T-001", name="Test", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.HIGH, impact=ImpactLevel.HIGH,
            mitigations=[Mitigation(name="Fix It", description="fix desc")],
        )
        card = _render_scenario_card(s)
        assert "Fix It" in card


class TestRenderRiskMatrix:
    def test_html_table(self) -> None:
        cells = [RiskMatrixCell(
            probability=Probability.HIGH,
            impact=ImpactLevel.HIGH,
            scenarios=["T-001"],
        )]
        html = _render_risk_matrix(cells)
        assert "<table" in html
        assert "T-001" in html


class TestRenderMitigationsSummary:
    def test_no_applicable(self) -> None:
        s = AttackScenario(
            id="T", name="T", category=AttackCategory.TOOL_ABUSE,
            description="d", vector="v",
            probability=Probability.LOW, impact=ImpactLevel.LOW,
            applicable=False,
        )
        html = _render_mitigations_summary([s])
        assert "No mitigations needed" in html

    def test_deduplicates(self) -> None:
        scenarios = []
        for i in range(3):
            scenarios.append(AttackScenario(
                id=f"T-{i}", name="T", category=AttackCategory.TOOL_ABUSE,
                description="d", vector="v",
                probability=Probability.HIGH, impact=ImpactLevel.HIGH,
                applicable=True,
                mitigations=[Mitigation(name="Same Fix", description="same")],
            ))
        html = _render_mitigations_summary(scenarios)
        assert html.count("Same Fix") == 1

    def test_sorts_by_priority(self) -> None:
        scenarios = [
            AttackScenario(
                id="T-1", name="T", category=AttackCategory.TOOL_ABUSE,
                description="d", vector="v",
                probability=Probability.LOW, impact=ImpactLevel.LOW,
                applicable=True,
                mitigations=[Mitigation(name="Low Fix", description="l", priority="LOW")],
            ),
            AttackScenario(
                id="T-2", name="T", category=AttackCategory.TOOL_ABUSE,
                description="d", vector="v",
                probability=Probability.HIGH, impact=ImpactLevel.HIGH,
                applicable=True,
                mitigations=[Mitigation(name="Critical Fix", description="c", priority="CRITICAL")],
            ),
        ]
        html = _render_mitigations_summary(scenarios)
        # Critical should appear before Low
        assert html.index("Critical Fix") < html.index("Low Fix")
