"""Tests for SAFE-MCP Threat Taxonomy module."""

import ast

from mcpaudit.models import Finding, Severity
from mcpaudit.taxonomy import (
    ALL_CATEGORIES,
    RULE_TAXONOMY,
    ThreatCategory,
    format_taxonomy_text,
    get_categories_for_rule,
    group_findings_by_category,
)


class TestThreatCategory:
    """Tests for ThreatCategory dataclass."""

    def test_category_to_dict(self):
        cat = ALL_CATEGORIES["confused_deputy"]
        d = cat.to_dict()
        assert d["id"] == "confused_deputy"
        assert d["name"] == "Confused Deputy"
        assert d["risk_level"] == 5
        assert "mitre_attack" in d

    def test_category_without_mitre(self):
        cat = ThreatCategory(
            id="test", name="Test", description="A test.", risk_level=1
        )
        d = cat.to_dict()
        assert "mitre_attack" not in d

    def test_all_categories_defined(self):
        expected = {
            "confused_deputy", "prompt_injection", "data_exfiltration",
            "privilege_escalation", "denial_of_service",
            "unauthorized_access", "information_disclosure", "supply_chain",
        }
        assert set(ALL_CATEGORIES.keys()) == expected

    def test_all_categories_have_descriptions(self):
        for cat_id, cat in ALL_CATEGORIES.items():
            assert cat.description, f"{cat_id} missing description"
            assert 1 <= cat.risk_level <= 5, f"{cat_id} invalid risk_level"


class TestRuleTaxonomyMapping:
    """Tests for rule-to-category mapping."""

    def test_path_001_maps_correctly(self):
        cats = get_categories_for_rule("PATH-001")
        ids = {c.id for c in cats}
        assert "data_exfiltration" in ids
        assert "unauthorized_access" in ids

    def test_cmd_001_maps_correctly(self):
        cats = get_categories_for_rule("CMD-001")
        ids = {c.id for c in cats}
        assert "privilege_escalation" in ids
        assert "confused_deputy" in ids

    def test_redos_maps_to_dos(self):
        cats = get_categories_for_rule("REDOS-001")
        ids = {c.id for c in cats}
        assert "denial_of_service" in ids

    def test_unknown_rule_returns_empty(self):
        cats = get_categories_for_rule("UNKNOWN-999")
        assert cats == []

    def test_mcp_003_mapped(self):
        cats = get_categories_for_rule("MCP-003")
        ids = {c.id for c in cats}
        assert "confused_deputy" in ids
        assert "privilege_escalation" in ids

    def test_mcp_004_mapped(self):
        cats = get_categories_for_rule("MCP-004")
        ids = {c.id for c in cats}
        assert "prompt_injection" in ids

    def test_all_mapped_categories_exist(self):
        for rule_id, cat_ids in RULE_TAXONOMY.items():
            for cat_id in cat_ids:
                assert cat_id in ALL_CATEGORIES, (
                    f"Rule {rule_id} maps to unknown category {cat_id}"
                )

    def test_multi_category_rules(self):
        multi = {
            rid for rid, cats in RULE_TAXONOMY.items() if len(cats) > 1
        }
        assert "CMD-001" in multi
        assert "PATH-001" in multi
        assert "SSRF-001" in multi


class TestGroupFindings:
    """Tests for grouping findings by threat category."""

    def _make_finding(self, rule_id: str, line: int = 1) -> Finding:
        return Finding(
            rule_id=rule_id,
            severity=Severity.HIGH,
            title=f"Test {rule_id}",
            description="test",
            file="test.py",
            line=line,
        )

    def test_single_finding_single_category(self):
        findings = [self._make_finding("REDOS-001")]
        grouped = group_findings_by_category(findings)
        assert "denial_of_service" in grouped
        assert len(grouped["denial_of_service"]) == 1

    def test_finding_appears_in_multiple_categories(self):
        findings = [self._make_finding("CMD-001")]
        grouped = group_findings_by_category(findings)
        assert "privilege_escalation" in grouped
        assert "confused_deputy" in grouped

    def test_empty_findings(self):
        grouped = group_findings_by_category([])
        assert grouped == {}

    def test_unmapped_finding_not_grouped(self):
        findings = [self._make_finding("UNKNOWN-999")]
        grouped = group_findings_by_category(findings)
        assert grouped == {}

    def test_multiple_findings_same_category(self):
        findings = [
            self._make_finding("REDOS-001", line=1),
            self._make_finding("RATE-001", line=5),
        ]
        grouped = group_findings_by_category(findings)
        assert len(grouped["denial_of_service"]) == 2


class TestFormatTaxonomyText:
    """Tests for taxonomy text output formatting."""

    def _make_finding(self, rule_id: str, line: int = 1) -> Finding:
        return Finding(
            rule_id=rule_id,
            severity=Severity.HIGH,
            title=f"Test {rule_id}",
            description="test finding",
            file="server.py",
            line=line,
        )

    def test_empty_findings_message(self):
        output = format_taxonomy_text([])
        assert "No findings to classify" in output

    def test_contains_category_header(self):
        findings = [self._make_finding("CMD-001")]
        output = format_taxonomy_text(findings)
        assert "Confused Deputy" in output
        assert "Privilege Escalation" in output

    def test_contains_finding_details(self):
        findings = [self._make_finding("CMD-001")]
        output = format_taxonomy_text(findings)
        assert "CMD-001" in output
        assert "server.py" in output

    def test_contains_mitre_reference(self):
        findings = [self._make_finding("REDOS-001")]
        output = format_taxonomy_text(findings)
        assert "MITRE ATT&CK" in output

    def test_contains_threat_taxonomy_header(self):
        findings = [self._make_finding("CMD-001")]
        output = format_taxonomy_text(findings)
        assert "SAFE-MCP Threat Taxonomy" in output

    def test_sorted_by_risk_level(self):
        findings = [
            self._make_finding("REDOS-001"),  # DoS, risk 3
            self._make_finding("CMD-001"),     # priv_esc risk 5
        ]
        output = format_taxonomy_text(findings)
        # Higher risk categories appear first
        pos_priv = output.find("Privilege Escalation")
        pos_dos = output.find("Denial of Service")
        assert pos_priv < pos_dos

    def test_summary_line(self):
        findings = [
            self._make_finding("CMD-001"),
            self._make_finding("REDOS-001"),
        ]
        output = format_taxonomy_text(findings)
        assert "threat categories" in output


class TestCLITaxonomyFlag:
    """Tests for --taxonomy CLI flag."""

    def test_taxonomy_flag_output(self, tmp_path, capsys):
        f = tmp_path / "test.py"
        f.write_text(
            'import subprocess\nsubprocess.run("ls", shell=True)\n'
        )
        from mcpaudit.__main__ import main
        main([str(f), "--taxonomy"])
        captured = capsys.readouterr()
        assert "SAFE-MCP Threat Taxonomy" in captured.out

    def test_taxonomy_flag_to_file(self, tmp_path):
        f = tmp_path / "test.py"
        f.write_text(
            'import subprocess\nsubprocess.run("ls", shell=True)\n'
        )
        out = tmp_path / "taxonomy.txt"
        from mcpaudit.__main__ import main
        main([str(f), "--taxonomy", "--output", str(out)])
        assert out.exists()
        content = out.read_text()
        assert "SAFE-MCP Threat Taxonomy" in content
