"""Tests for the Threat Model Generator module."""

from __future__ import annotations

import json

import pytest

from mcpaudit.threat_model import (
    Asset,
    AssetType,
    AttackVector,
    Impact,
    Likelihood,
    Mitigation,
    StrideCategory,
    Threat,
    ThreatAgent,
    ThreatModel,
    format_threat_model_html,
    format_threat_model_json,
    format_threat_model_text,
    generate_threat_model,
    _build_mitigations,
    _classify_rule,
    _default_threat_agents,
    _escape_html,
    _extract_assets_from_scan,
)


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

def _scan_data(
    grade: str = "C",
    findings: list[dict] | None = None,
    total: int = 5,
) -> dict:
    """Build a minimal ScanSummary.to_dict()-compatible dict."""
    if findings is None:
        findings = [
            {
                "rule_id": "CMD-001",
                "severity": "CRITICAL",
                "title": "Command Injection",
                "description": "Unsanitized input in subprocess call",
                "file": "server.py",
                "line": 42,
                "recommendation": "Use shlex.quote() to sanitize input",
            },
            {
                "rule_id": "AUTH-001",
                "severity": "HIGH",
                "title": "Missing Authentication",
                "description": "No authentication on tool endpoint",
                "file": "server.py",
                "line": 10,
                "recommendation": "Add authentication middleware",
            },
            {
                "rule_id": "RATE-001",
                "severity": "MEDIUM",
                "title": "No Rate Limiting",
                "description": "Tool endpoint has no rate limiting",
                "file": "tools.py",
                "line": 25,
                "recommendation": "Add rate limiting middleware",
            },
        ]
    return {
        "summary": {
            "files_scanned": 2,
            "total_findings": total,
            "overall_grade": grade,
            "severity_counts": {"CRITICAL": 1, "HIGH": 1, "MEDIUM": 1, "LOW": 0, "INFO": 0},
        },
        "results": [
            {
                "file": "server.py",
                "findings": findings[:2] if len(findings) >= 2 else findings,
            },
            {
                "file": "tools.py",
                "findings": findings[2:] if len(findings) > 2 else [],
            },
        ],
    }


# ===========================================================================
# StrideCategory
# ===========================================================================

class TestStrideCategory:
    def test_all_values(self):
        assert len(StrideCategory) == 6

    def test_from_string_basic(self):
        assert StrideCategory.from_string("spoofing") == StrideCategory.SPOOFING
        assert StrideCategory.from_string("tampering") == StrideCategory.TAMPERING
        assert StrideCategory.from_string("repudiation") == StrideCategory.REPUDIATION

    def test_from_string_info_disclosure(self):
        assert StrideCategory.from_string("info_disclosure") == StrideCategory.INFO_DISCLOSURE
        assert StrideCategory.from_string("information disclosure") == StrideCategory.INFO_DISCLOSURE
        assert StrideCategory.from_string("information_disclosure") == StrideCategory.INFO_DISCLOSURE

    def test_from_string_dos(self):
        assert StrideCategory.from_string("dos") == StrideCategory.DENIAL_OF_SERVICE
        assert StrideCategory.from_string("denial_of_service") == StrideCategory.DENIAL_OF_SERVICE
        assert StrideCategory.from_string("denial of service") == StrideCategory.DENIAL_OF_SERVICE

    def test_from_string_elevation(self):
        assert StrideCategory.from_string("elevation") == StrideCategory.ELEVATION
        assert StrideCategory.from_string("elevation_of_privilege") == StrideCategory.ELEVATION
        assert StrideCategory.from_string("elevation of privilege") == StrideCategory.ELEVATION

    def test_from_string_case_insensitive(self):
        assert StrideCategory.from_string("SPOOFING") == StrideCategory.SPOOFING
        assert StrideCategory.from_string("Tampering") == StrideCategory.TAMPERING

    def test_from_string_invalid(self):
        with pytest.raises(ValueError, match="Unknown STRIDE"):
            StrideCategory.from_string("nonexistent")

    def test_values(self):
        assert StrideCategory.SPOOFING.value == "Spoofing"
        assert StrideCategory.ELEVATION.value == "Elevation of Privilege"


# ===========================================================================
# Likelihood
# ===========================================================================

class TestLikelihood:
    def test_all_values(self):
        assert len(Likelihood) == 5

    def test_scores(self):
        assert Likelihood.VERY_HIGH.score == 5
        assert Likelihood.HIGH.score == 4
        assert Likelihood.MEDIUM.score == 3
        assert Likelihood.LOW.score == 2
        assert Likelihood.VERY_LOW.score == 1

    def test_values(self):
        assert Likelihood.VERY_HIGH.value == "Very High"


# ===========================================================================
# Impact
# ===========================================================================

class TestImpact:
    def test_all_values(self):
        assert len(Impact) == 5

    def test_scores(self):
        assert Impact.CRITICAL.score == 5
        assert Impact.HIGH.score == 4
        assert Impact.MEDIUM.score == 3
        assert Impact.LOW.score == 2
        assert Impact.NEGLIGIBLE.score == 1

    def test_values(self):
        assert Impact.CRITICAL.value == "Critical"


# ===========================================================================
# AssetType
# ===========================================================================

class TestAssetType:
    def test_all_values(self):
        assert len(AssetType) == 7
        assert AssetType.TOOL.value == "Tool"
        assert AssetType.DATA.value == "Data"
        assert AssetType.CONFIG.value == "Configuration"
        assert AssetType.CREDENTIAL.value == "Credential"
        assert AssetType.RESOURCE.value == "Resource"
        assert AssetType.PROMPT.value == "Prompt"
        assert AssetType.ENDPOINT.value == "Endpoint"


# ===========================================================================
# Asset
# ===========================================================================

class TestAsset:
    def test_to_dict(self):
        a = Asset("myfile.py", AssetType.TOOL, "A tool file", "high")
        d = a.to_dict()
        assert d["name"] == "myfile.py"
        assert d["type"] == "Tool"
        assert d["sensitivity"] == "high"

    def test_defaults(self):
        a = Asset("x", AssetType.DATA, "desc")
        assert a.sensitivity == "medium"


# ===========================================================================
# ThreatAgent
# ===========================================================================

class TestThreatAgent:
    def test_to_dict(self):
        ta = ThreatAgent("Attacker", "A bad actor", "high", "high")
        d = ta.to_dict()
        assert d["name"] == "Attacker"
        assert d["capability"] == "high"

    def test_defaults(self):
        ta = ThreatAgent("X", "desc")
        assert ta.capability == "medium"
        assert ta.motivation == "medium"


# ===========================================================================
# AttackVector
# ===========================================================================

class TestAttackVector:
    def test_to_dict(self):
        av = AttackVector("SQLi", "SQL injection", "T1190", ["network access"])
        d = av.to_dict()
        assert d["name"] == "SQLi"
        assert d["technique"] == "T1190"
        assert d["prerequisites"] == ["network access"]

    def test_defaults(self):
        av = AttackVector("x", "desc")
        assert av.technique == ""
        assert av.prerequisites == []


# ===========================================================================
# Mitigation
# ===========================================================================

class TestMitigation:
    def test_to_dict(self):
        m = Mitigation("Fix auth", "Add auth", "critical", "implemented", "low")
        d = m.to_dict()
        assert d["name"] == "Fix auth"
        assert d["priority"] == "critical"
        assert d["status"] == "implemented"
        assert d["effort"] == "low"

    def test_defaults(self):
        m = Mitigation("x", "desc")
        assert m.priority == "medium"
        assert m.status == "recommended"
        assert m.effort == "medium"


# ===========================================================================
# Threat
# ===========================================================================

class TestThreat:
    def _make_threat(self, lh=Likelihood.HIGH, imp=Impact.HIGH):
        return Threat(
            threat_id="T-001",
            name="Test Threat",
            description="A test threat",
            stride_category=StrideCategory.TAMPERING,
            likelihood=lh,
            impact=imp,
        )

    def test_risk_score(self):
        t = self._make_threat(Likelihood.HIGH, Impact.CRITICAL)
        assert t.risk_score == 20  # 4 * 5

    def test_risk_level_critical(self):
        t = self._make_threat(Likelihood.VERY_HIGH, Impact.CRITICAL)
        assert t.risk_level == "critical"  # 25

    def test_risk_level_high(self):
        t = self._make_threat(Likelihood.HIGH, Impact.HIGH)
        assert t.risk_level == "high"  # 16

    def test_risk_level_medium(self):
        t = self._make_threat(Likelihood.MEDIUM, Impact.MEDIUM)
        assert t.risk_level == "medium"  # 9

    def test_risk_level_low(self):
        t = self._make_threat(Likelihood.VERY_LOW, Impact.LOW)
        assert t.risk_level == "low"  # 2

    def test_to_dict(self):
        t = self._make_threat()
        t.mitigations.append(Mitigation("m1", "desc"))
        t.attack_vectors.append(AttackVector("av1", "desc"))
        d = t.to_dict()
        assert d["threat_id"] == "T-001"
        assert d["stride_category"] == "Tampering"
        assert isinstance(d["mitigations"], list)
        assert isinstance(d["attack_vectors"], list)
        assert "risk_score" in d
        assert "risk_level" in d

    def test_defaults(self):
        t = self._make_threat()
        assert t.affected_assets == []
        assert t.attack_vectors == []
        assert t.mitigations == []
        assert t.related_findings == []


# ===========================================================================
# ThreatModel
# ===========================================================================

class TestThreatModel:
    def _make_model(self, n_threats=3):
        model = ThreatModel(name="Test", description="desc")
        categories = [StrideCategory.SPOOFING, StrideCategory.TAMPERING, StrideCategory.ELEVATION]
        for i in range(n_threats):
            model.threats.append(Threat(
                threat_id=f"T-{i+1:03d}",
                name=f"Threat {i+1}",
                description=f"desc {i+1}",
                stride_category=categories[i % len(categories)],
                likelihood=Likelihood.HIGH,
                impact=Impact.HIGH if i < 2 else Impact.MEDIUM,
            ))
        return model

    def test_threat_count(self):
        m = self._make_model(3)
        assert m.threat_count == 3

    def test_stride_summary(self):
        m = self._make_model(3)
        s = m.stride_summary
        assert s["Spoofing"] == 1
        assert s["Tampering"] == 1
        assert s["Elevation of Privilege"] == 1

    def test_risk_matrix(self):
        m = self._make_model(2)
        matrix = m.risk_matrix
        assert len(matrix) == 2
        assert "threat_id" in matrix[0]
        assert "risk_score" in matrix[0]
        # Should be sorted by risk score descending
        assert matrix[0]["risk_score"] >= matrix[1]["risk_score"]

    def test_critical_threats(self):
        m = ThreatModel(name="T", description="d")
        m.threats.append(Threat(
            "T-1", "Critical", "d", StrideCategory.ELEVATION,
            Likelihood.VERY_HIGH, Impact.CRITICAL,
        ))
        m.threats.append(Threat(
            "T-2", "Low", "d", StrideCategory.SPOOFING,
            Likelihood.LOW, Impact.LOW,
        ))
        assert len(m.critical_threats) == 1
        assert m.critical_threats[0].threat_id == "T-1"

    def test_high_threats(self):
        m = ThreatModel(name="T", description="d")
        m.threats.append(Threat(
            "T-1", "High", "d", StrideCategory.TAMPERING,
            Likelihood.HIGH, Impact.HIGH,
        ))
        assert len(m.high_threats) == 1

    def test_to_dict(self):
        m = self._make_model(1)
        m.assets.append(Asset("f", AssetType.TOOL, "d"))
        m.threat_agents.append(ThreatAgent("TA", "d"))
        d = m.to_dict()
        assert d["name"] == "Test"
        assert "stride_summary" in d
        assert "risk_matrix" in d
        assert len(d["assets"]) == 1
        assert len(d["threat_agents"]) == 1

    def test_empty_model(self):
        m = ThreatModel(name="Empty", description="empty")
        assert m.threat_count == 0
        assert m.critical_threats == []
        assert m.high_threats == []
        assert m.stride_summary == {c.value: 0 for c in StrideCategory}


# ===========================================================================
# Helper functions
# ===========================================================================

class TestClassifyRule:
    def test_cmd(self):
        assert _classify_rule("CMD-001") == "CMD"

    def test_auth(self):
        assert _classify_rule("AUTH-002") == "AUTH"

    def test_prompt(self):
        assert _classify_rule("PROMPT-001") == "PROMPT"

    def test_unknown(self):
        assert _classify_rule("UNKNOWN-001") == ""

    def test_partial_match(self):
        assert _classify_rule("CONFUSED-DEPUTY-001") == "CONFUSED"


class TestDefaultThreatAgents:
    def test_returns_list(self):
        agents = _default_threat_agents()
        assert isinstance(agents, list)
        assert len(agents) >= 4

    def test_agent_properties(self):
        agents = _default_threat_agents()
        for a in agents:
            assert a.name
            assert a.description
            assert a.capability in ("low", "medium", "high")
            assert a.motivation in ("low", "medium", "high")


class TestExtractAssets:
    def test_basic_extraction(self):
        data = _scan_data()
        assets = _extract_assets_from_scan(data)
        assert len(assets) >= 2  # At least files + transport layer

    def test_adds_transport_layer(self):
        data = _scan_data()
        assets = _extract_assets_from_scan(data)
        names = [a.name for a in assets]
        assert "MCP Transport Layer" in names

    def test_config_file_detection(self):
        data = {
            "results": [
                {"file": "config.json", "findings": []},
            ],
        }
        assets = _extract_assets_from_scan(data)
        config_assets = [a for a in assets if a.asset_type == AssetType.CONFIG]
        assert len(config_assets) >= 1

    def test_empty_scan(self):
        assets = _extract_assets_from_scan({"results": []})
        # Should still have transport layer
        assert len(assets) >= 1

    def test_no_duplicate_files(self):
        data = {
            "results": [
                {"file": "server.py", "findings": [{"rule_id": "CMD-001"}]},
                {"file": "server.py", "findings": [{"rule_id": "AUTH-001"}]},
            ],
        }
        assets = _extract_assets_from_scan(data)
        file_assets = [a for a in assets if a.name == "server.py"]
        assert len(file_assets) == 1


class TestBuildMitigations:
    def test_spoofing(self):
        ms = _build_mitigations(StrideCategory.SPOOFING, "AUTH-001", "Fix auth")
        assert len(ms) >= 2  # Recommendation + generic
        assert ms[0].name == "Fix AUTH-001"

    def test_tampering(self):
        ms = _build_mitigations(StrideCategory.TAMPERING, "INJ-001", "")
        assert any("input validation" in m.name.lower() for m in ms)

    def test_repudiation(self):
        ms = _build_mitigations(StrideCategory.REPUDIATION, "LOG-001", "Add logging")
        assert len(ms) >= 2

    def test_info_disclosure(self):
        ms = _build_mitigations(StrideCategory.INFO_DISCLOSURE, "LEAK-001", "")
        assert any("encryption" in m.name.lower() for m in ms)

    def test_dos(self):
        ms = _build_mitigations(StrideCategory.DENIAL_OF_SERVICE, "RATE-001", "")
        assert any("rate" in m.name.lower() for m in ms)

    def test_elevation(self):
        ms = _build_mitigations(StrideCategory.ELEVATION, "CMD-001", "Fix it")
        assert any("privilege" in m.name.lower() or "sandbox" in m.name.lower() for m in ms)

    def test_no_duplicate_names(self):
        ms = _build_mitigations(StrideCategory.SPOOFING, "AUTH-001", "Implement authentication")
        names = [m.name for m in ms]
        # The recommendation might overlap with generic name
        assert len(ms) >= 2


class TestEscapeHtml:
    def test_ampersand(self):
        assert _escape_html("a & b") == "a &amp; b"

    def test_lt_gt(self):
        assert _escape_html("<script>") == "&lt;script&gt;"

    def test_quote(self):
        assert _escape_html('"hello"') == "&quot;hello&quot;"

    def test_no_change(self):
        assert _escape_html("plain text") == "plain text"


# ===========================================================================
# generate_threat_model
# ===========================================================================

class TestGenerateThreatModel:
    def test_basic_generation(self):
        data = _scan_data()
        model = generate_threat_model(data)
        assert isinstance(model, ThreatModel)
        assert model.threat_count >= 3  # CMD, AUTH, RATE
        assert model.scan_grade == "C"
        assert model.total_findings == 5

    def test_empty_scan(self):
        data = {"summary": {}, "results": []}
        model = generate_threat_model(data)
        assert model.threat_count == 0
        assert model.scan_grade == "?"

    def test_assets_populated(self):
        data = _scan_data()
        model = generate_threat_model(data)
        assert len(model.assets) >= 1

    def test_threat_agents_populated(self):
        data = _scan_data()
        model = generate_threat_model(data)
        assert len(model.threat_agents) >= 4

    def test_threats_have_mitigations(self):
        data = _scan_data()
        model = generate_threat_model(data)
        for t in model.threats:
            assert len(t.mitigations) >= 1

    def test_threats_have_attack_vectors(self):
        data = _scan_data()
        model = generate_threat_model(data)
        for t in model.threats:
            assert len(t.attack_vectors) >= 1

    def test_stride_categories_assigned(self):
        data = _scan_data()
        model = generate_threat_model(data)
        categories = {t.stride_category for t in model.threats}
        assert len(categories) >= 2  # Multiple categories

    def test_severity_boost(self):
        data = _scan_data(findings=[
            {
                "rule_id": "CMD-001",
                "severity": "CRITICAL",
                "title": "Critical cmd",
                "description": "desc",
                "file": "f.py",
                "line": 1,
                "recommendation": "fix",
            },
        ])
        model = generate_threat_model(data)
        critical_threats = [t for t in model.threats if t.impact == Impact.CRITICAL]
        assert len(critical_threats) >= 1

    def test_related_findings_tracked(self):
        data = _scan_data()
        model = generate_threat_model(data)
        for t in model.threats:
            assert len(t.related_findings) >= 1

    def test_risk_matrix_sorted(self):
        data = _scan_data()
        model = generate_threat_model(data)
        matrix = model.risk_matrix
        if len(matrix) >= 2:
            for i in range(len(matrix) - 1):
                assert matrix[i]["risk_score"] >= matrix[i + 1]["risk_score"]


# ===========================================================================
# format_threat_model_text
# ===========================================================================

class TestFormatThreatModelText:
    def test_basic_output(self):
        data = _scan_data()
        model = generate_threat_model(data)
        text = format_threat_model_text(model)
        assert "MCP Server Threat Model" in text
        assert "STRIDE Summary" in text

    def test_contains_grade(self):
        data = _scan_data(grade="B")
        model = generate_threat_model(data)
        text = format_threat_model_text(model)
        assert "Grade: B" in text

    def test_contains_risk_matrix(self):
        data = _scan_data()
        model = generate_threat_model(data)
        text = format_threat_model_text(model)
        assert "Likelihood" in text
        assert "Impact" in text

    def test_contains_threat_details(self):
        data = _scan_data()
        model = generate_threat_model(data)
        text = format_threat_model_text(model)
        assert "T-001" in text
        assert "Mitigations" in text

    def test_empty_model(self):
        model = ThreatModel(name="Empty", description="empty")
        text = format_threat_model_text(model)
        assert "Empty" in text
        assert "Threats Identified: 0" in text


# ===========================================================================
# format_threat_model_json
# ===========================================================================

class TestFormatThreatModelJson:
    def test_valid_json(self):
        data = _scan_data()
        model = generate_threat_model(data)
        j = format_threat_model_json(model)
        parsed = json.loads(j)
        assert "threats" in parsed
        assert "risk_matrix" in parsed

    def test_round_trip(self):
        data = _scan_data()
        model = generate_threat_model(data)
        j = format_threat_model_json(model)
        parsed = json.loads(j)
        assert parsed["name"] == model.name
        assert parsed["threat_count"] == model.threat_count

    def test_empty_model(self):
        model = ThreatModel(name="E", description="e")
        j = format_threat_model_json(model)
        parsed = json.loads(j)
        assert parsed["threats"] == []


# ===========================================================================
# format_threat_model_html
# ===========================================================================

class TestFormatThreatModelHtml:
    def test_valid_html(self):
        data = _scan_data()
        model = generate_threat_model(data)
        html = format_threat_model_html(model)
        assert "<!DOCTYPE html>" in html
        assert "</html>" in html

    def test_contains_title(self):
        data = _scan_data()
        model = generate_threat_model(data)
        html = format_threat_model_html(model)
        assert "MCP Threat Model" in html

    def test_contains_stride_summary(self):
        data = _scan_data()
        model = generate_threat_model(data)
        html = format_threat_model_html(model)
        assert "STRIDE Summary" in html
        assert "stride-card" in html

    def test_contains_risk_matrix(self):
        data = _scan_data()
        model = generate_threat_model(data)
        html = format_threat_model_html(model)
        assert "Risk Matrix" in html
        assert "T-001" in html

    def test_contains_threat_cards(self):
        data = _scan_data()
        model = generate_threat_model(data)
        html = format_threat_model_html(model)
        assert "threat-card" in html
        assert "Mitigations" in html
        assert "Attack Vectors" in html

    def test_contains_assets(self):
        data = _scan_data()
        model = generate_threat_model(data)
        html = format_threat_model_html(model)
        assert "Assets" in html

    def test_escapes_html(self):
        data = _scan_data(findings=[
            {
                "rule_id": "CMD-001",
                "severity": "HIGH",
                "title": "Test <script>",
                "description": "A & B",
                "file": "f.py",
                "line": 1,
                "recommendation": "fix",
            },
        ])
        model = generate_threat_model(data)
        html = format_threat_model_html(model)
        assert "<script>" not in html
        assert "&lt;script&gt;" in html

    def test_empty_model(self):
        model = ThreatModel(name="E", description="e")
        html = format_threat_model_html(model)
        assert "<!DOCTYPE html>" in html
        assert "0" in html
