"""Tests for MCP Threat Intelligence Database."""

from __future__ import annotations

import json

import pytest

from mcpaudit.threat_intel import (
    ThreatEntry,
    ThreatIntelDB,
    ThreatMatch,
    ThreatReport,
    ThreatSeverity,
    _build_threat_db,
    _compute_confidence,
    format_threat_json,
    format_threat_text,
    format_threats_list_text,
    generate_threat_report,
    match_threats,
    match_threats_from_scan,
)


# --- ThreatSeverity ---


class TestThreatSeverity:
    def test_all_values(self):
        assert ThreatSeverity.CRITICAL.value == "CRITICAL"
        assert ThreatSeverity.HIGH.value == "HIGH"
        assert ThreatSeverity.MEDIUM.value == "MEDIUM"
        assert ThreatSeverity.LOW.value == "LOW"
        assert ThreatSeverity.INFO.value == "INFO"

    def test_from_string_valid(self):
        assert ThreatSeverity.from_string("critical") == ThreatSeverity.CRITICAL
        assert ThreatSeverity.from_string("HIGH") == ThreatSeverity.HIGH
        assert ThreatSeverity.from_string("Medium") == ThreatSeverity.MEDIUM

    def test_from_string_invalid(self):
        with pytest.raises(ValueError, match="Invalid severity"):
            ThreatSeverity.from_string("unknown")

    def test_enum_iteration(self):
        all_sevs = list(ThreatSeverity)
        assert len(all_sevs) == 5


# --- ThreatEntry ---


class TestThreatEntry:
    def test_creation(self):
        entry = ThreatEntry(
            id="TEST-001",
            name="Test Threat",
            description="A test threat",
            severity=ThreatSeverity.HIGH,
        )
        assert entry.id == "TEST-001"
        assert entry.name == "Test Threat"
        assert entry.severity == ThreatSeverity.HIGH

    def test_to_dict(self):
        entry = ThreatEntry(
            id="TEST-001",
            name="Test",
            description="Desc",
            severity=ThreatSeverity.MEDIUM,
            tags=["test"],
            rule_ids=["CMD-001"],
        )
        d = entry.to_dict()
        assert d["id"] == "TEST-001"
        assert d["severity"] == "MEDIUM"
        assert d["tags"] == ["test"]
        assert d["rule_ids"] == ["CMD-001"]

    def test_from_dict(self):
        data = {
            "id": "TEST-002",
            "name": "From Dict",
            "description": "Created from dict",
            "severity": "HIGH",
            "tags": ["tag1"],
            "rule_ids": ["INJ-001"],
            "affected_tools": ["tool_a"],
            "indicators": ["indicator_1"],
            "mitigation": "Fix it",
            "references": ["ref1"],
        }
        entry = ThreatEntry.from_dict(data)
        assert entry.id == "TEST-002"
        assert entry.severity == ThreatSeverity.HIGH
        assert entry.tags == ["tag1"]
        assert entry.affected_tools == ["tool_a"]

    def test_from_dict_minimal(self):
        data = {
            "id": "MIN-001",
            "name": "Minimal",
            "description": "Minimal entry",
            "severity": "LOW",
        }
        entry = ThreatEntry.from_dict(data)
        assert entry.affected_tools == []
        assert entry.indicators == []
        assert entry.mitigation == ""

    def test_roundtrip(self):
        entry = ThreatEntry(
            id="RT-001",
            name="Roundtrip",
            description="Test roundtrip",
            severity=ThreatSeverity.CRITICAL,
            tags=["a", "b"],
            rule_ids=["X-001"],
        )
        restored = ThreatEntry.from_dict(entry.to_dict())
        assert restored.id == entry.id
        assert restored.severity == entry.severity
        assert restored.tags == entry.tags

    def test_default_fields(self):
        entry = ThreatEntry(
            id="DEF-001",
            name="Defaults",
            description="Defaults test",
            severity=ThreatSeverity.INFO,
        )
        assert entry.affected_tools == []
        assert entry.indicators == []
        assert entry.mitigation == ""
        assert entry.references == []
        assert entry.tags == []
        assert entry.rule_ids == []


# --- ThreatMatch ---


class TestThreatMatch:
    def test_creation(self):
        threat = ThreatEntry(
            id="T-001", name="T", description="D",
            severity=ThreatSeverity.HIGH,
        )
        match = ThreatMatch(
            threat=threat,
            finding_rule_id="CMD-001",
            finding_title="Command injection",
            finding_file="server.py",
            finding_line=42,
        )
        assert match.finding_rule_id == "CMD-001"
        assert match.confidence == "HIGH"

    def test_to_dict(self):
        threat = ThreatEntry(
            id="T-002", name="Threat", description="D",
            severity=ThreatSeverity.CRITICAL,
            mitigation="Fix it",
        )
        match = ThreatMatch(
            threat=threat,
            finding_rule_id="INJ-001",
            finding_title="Injection",
            finding_file="app.py",
            finding_line=10,
            confidence="MEDIUM",
        )
        d = match.to_dict()
        assert d["threat_id"] == "T-002"
        assert d["threat_severity"] == "CRITICAL"
        assert d["confidence"] == "MEDIUM"
        assert d["mitigation"] == "Fix it"


# --- ThreatReport ---


class TestThreatReport:
    def test_empty_report(self):
        report = ThreatReport()
        assert report.total_threats == 0
        assert report.total_findings_analyzed == 0
        assert len(report.matches) == 0

    def test_severity_counts(self):
        t1 = ThreatEntry(
            id="A", name="A", description="A",
            severity=ThreatSeverity.CRITICAL,
        )
        t2 = ThreatEntry(
            id="B", name="B", description="B",
            severity=ThreatSeverity.HIGH,
        )
        report = ThreatReport(
            matches=[
                ThreatMatch(t1, "R1", "T1", "f.py", 1),
                ThreatMatch(t1, "R2", "T2", "f.py", 2),
                ThreatMatch(t2, "R3", "T3", "f.py", 3),
            ],
            total_threats=5,
        )
        counts = report.severity_counts
        assert counts["CRITICAL"] == 2
        assert counts["HIGH"] == 1
        assert report.critical_count == 2
        assert report.high_count == 1

    def test_to_dict(self):
        report = ThreatReport(total_threats=10, total_findings_analyzed=5)
        d = report.to_dict()
        assert d["total_threats_in_db"] == 10
        assert d["matches_found"] == 0

    def test_unmatched_threats(self):
        report = ThreatReport(
            unmatched_threats=["T-001", "T-002"],
        )
        assert len(report.unmatched_threats) == 2


# --- Built-in Threat DB ---


class TestBuiltinDB:
    def test_has_minimum_threats(self):
        threats = _build_threat_db()
        assert len(threats) >= 20

    def test_all_have_ids(self):
        for t in _build_threat_db():
            assert t.id.startswith("THREAT-")

    def test_all_have_descriptions(self):
        for t in _build_threat_db():
            assert len(t.description) > 10

    def test_all_have_mitigation(self):
        for t in _build_threat_db():
            assert len(t.mitigation) > 10

    def test_all_have_rule_ids(self):
        for t in _build_threat_db():
            assert len(t.rule_ids) > 0

    def test_unique_ids(self):
        threats = _build_threat_db()
        ids = [t.id for t in threats]
        assert len(ids) == len(set(ids))

    def test_severity_distribution(self):
        threats = _build_threat_db()
        severities = {t.severity for t in threats}
        # At least CRITICAL, HIGH, and MEDIUM should be represented
        assert ThreatSeverity.CRITICAL in severities
        assert ThreatSeverity.HIGH in severities
        assert ThreatSeverity.MEDIUM in severities


# --- ThreatIntelDB ---


class TestThreatIntelDB:
    def test_default_init(self):
        db = ThreatIntelDB()
        assert db.count >= 20

    def test_custom_threats(self):
        threats = [
            ThreatEntry(
                id="CUSTOM-001", name="Custom", description="Custom threat",
                severity=ThreatSeverity.LOW, rule_ids=["X-001"],
            ),
        ]
        db = ThreatIntelDB(threats)
        assert db.count == 1

    def test_get_threat(self):
        db = ThreatIntelDB()
        t = db.get_threat("THREAT-001")
        assert t is not None
        assert t.name == "Confused Deputy Attack"

    def test_get_threat_missing(self):
        db = ThreatIntelDB()
        assert db.get_threat("NONEXISTENT") is None

    def test_list_threats(self):
        db = ThreatIntelDB()
        all_threats = db.list_threats()
        assert len(all_threats) == db.count

    def test_filter_by_severity(self):
        db = ThreatIntelDB()
        critical = db.filter_by_severity(ThreatSeverity.CRITICAL)
        assert len(critical) > 0
        assert all(t.severity == ThreatSeverity.CRITICAL for t in critical)

    def test_filter_by_tag(self):
        db = ThreatIntelDB()
        owasp = db.filter_by_tag("owasp")
        assert len(owasp) > 0
        for t in owasp:
            assert "owasp" in [tag.lower() for tag in t.tags]

    def test_filter_by_tag_case_insensitive(self):
        db = ThreatIntelDB()
        r1 = db.filter_by_tag("OWASP")
        r2 = db.filter_by_tag("owasp")
        assert len(r1) == len(r2)

    def test_search_by_name(self):
        db = ThreatIntelDB()
        results = db.search("injection")
        assert len(results) > 0

    def test_search_by_id(self):
        db = ThreatIntelDB()
        results = db.search("THREAT-001")
        assert len(results) == 1

    def test_search_no_results(self):
        db = ThreatIntelDB()
        results = db.search("xyznonexistent123")
        assert len(results) == 0

    def test_threats_for_rule(self):
        db = ThreatIntelDB()
        threats = db.threats_for_rule("CMD-001")
        assert len(threats) > 0

    def test_threats_for_unknown_rule(self):
        db = ThreatIntelDB()
        threats = db.threats_for_rule("UNKNOWN-999")
        assert len(threats) == 0


# --- match_threats ---


class TestMatchThreats:
    def test_match_cmd_finding(self):
        findings = [
            {"rule_id": "CMD-001", "title": "Command injection", "file": "s.py", "line": 10},
        ]
        matches = match_threats(findings)
        assert len(matches) > 0
        assert any(m.finding_rule_id == "CMD-001" for m in matches)

    def test_no_match_for_unknown_rule(self):
        findings = [
            {"rule_id": "UNKNOWN-999", "title": "Unknown", "file": "x.py", "line": 1},
        ]
        matches = match_threats(findings)
        assert len(matches) == 0

    def test_multiple_findings(self):
        findings = [
            {"rule_id": "CMD-001", "title": "Cmd inj", "file": "a.py", "line": 1},
            {"rule_id": "INJ-001", "title": "Prompt inj", "file": "b.py", "line": 2},
        ]
        matches = match_threats(findings)
        assert len(matches) >= 2

    def test_custom_db(self):
        db = ThreatIntelDB([
            ThreatEntry(
                id="CUSTOM-001", name="Custom", description="Custom",
                severity=ThreatSeverity.LOW, rule_ids=["CUSTOM-R1"],
            ),
        ])
        findings = [
            {"rule_id": "CUSTOM-R1", "title": "Custom", "file": "c.py", "line": 5},
        ]
        matches = match_threats(findings, db)
        assert len(matches) == 1
        assert matches[0].threat.id == "CUSTOM-001"

    def test_empty_findings(self):
        matches = match_threats([])
        assert len(matches) == 0

    def test_finding_missing_fields(self):
        findings = [{"rule_id": "CMD-001"}]
        matches = match_threats(findings)
        assert len(matches) > 0
        assert matches[0].finding_title == ""

    def test_match_returns_correct_file_info(self):
        findings = [
            {"rule_id": "CMD-001", "title": "Inj", "file": "my_server.py", "line": 99},
        ]
        matches = match_threats(findings)
        assert any(
            m.finding_file == "my_server.py" and m.finding_line == 99
            for m in matches
        )


# --- match_threats_from_scan ---


class TestMatchThreatsFromScan:
    def test_scan_result_format(self):
        scan_result = {
            "results": [
                {
                    "file": "server.py",
                    "findings": [
                        {"rule_id": "CMD-001", "title": "CI", "file": "server.py", "line": 5},
                    ],
                },
            ],
        }
        matches = match_threats_from_scan(scan_result)
        assert len(matches) > 0

    def test_empty_scan(self):
        matches = match_threats_from_scan({"results": []})
        assert len(matches) == 0

    def test_scan_no_results_key(self):
        matches = match_threats_from_scan({})
        assert len(matches) == 0

    def test_multiple_files(self):
        scan_result = {
            "results": [
                {
                    "file": "a.py",
                    "findings": [
                        {"rule_id": "CMD-001", "title": "A", "file": "a.py", "line": 1},
                    ],
                },
                {
                    "file": "b.py",
                    "findings": [
                        {"rule_id": "INJ-001", "title": "B", "file": "b.py", "line": 2},
                    ],
                },
            ],
        }
        matches = match_threats_from_scan(scan_result)
        assert len(matches) >= 2


# --- _compute_confidence ---


class TestComputeConfidence:
    def test_direct_rule_match(self):
        threat = ThreatEntry(
            id="T1", name="T", description="D",
            severity=ThreatSeverity.HIGH,
            rule_ids=["CMD-001"],
            indicators=["subprocess.call with shell=True"],
        )
        conf = _compute_confidence(
            {"rule_id": "CMD-001", "title": "subprocess shell", "description": ""},
            threat,
        )
        assert conf in ("HIGH", "MEDIUM")

    def test_no_rule_match(self):
        threat = ThreatEntry(
            id="T1", name="T", description="D",
            severity=ThreatSeverity.HIGH,
            rule_ids=["OTHER-001"],
        )
        conf = _compute_confidence(
            {"rule_id": "CMD-001", "title": "test", "description": ""},
            threat,
        )
        assert conf == "LOW"


# --- generate_threat_report ---


class TestGenerateThreatReport:
    def test_basic_report(self):
        findings = [
            {"rule_id": "CMD-001", "title": "Cmd", "file": "s.py", "line": 1},
        ]
        matches = match_threats(findings)
        report = generate_threat_report(matches)
        assert report.total_threats >= 20
        assert len(report.matches) > 0

    def test_empty_report(self):
        report = generate_threat_report([])
        assert report.total_threats >= 20
        assert len(report.matches) == 0
        assert len(report.unmatched_threats) == report.total_threats

    def test_unmatched_threats(self):
        findings = [
            {"rule_id": "CMD-001", "title": "Cmd", "file": "s.py", "line": 1},
        ]
        matches = match_threats(findings)
        report = generate_threat_report(matches)
        matched_ids = {m.threat.id for m in matches}
        for uid in report.unmatched_threats:
            assert uid not in matched_ids


# --- Formatting ---


class TestFormatting:
    def test_format_text_empty(self):
        report = ThreatReport(total_threats=10)
        text = format_threat_text(report)
        assert "Threat Intelligence Report" in text
        assert "Threats in database: 10" in text

    def test_format_text_with_matches(self):
        threat = ThreatEntry(
            id="T-001", name="Test Threat", description="A test description here",
            severity=ThreatSeverity.HIGH,
            mitigation="Apply the fix",
        )
        report = ThreatReport(
            matches=[
                ThreatMatch(threat, "R1", "Title", "f.py", 1),
            ],
            total_threats=1,
        )
        text = format_threat_text(report)
        assert "T-001" in text
        assert "Test Threat" in text

    def test_format_json(self):
        report = ThreatReport(total_threats=5)
        j = format_threat_json(report)
        data = json.loads(j)
        assert data["total_threats_in_db"] == 5
        assert data["matches_found"] == 0

    def test_format_json_roundtrip(self):
        findings = [
            {"rule_id": "CMD-001", "title": "Cmd", "file": "s.py", "line": 1},
        ]
        matches = match_threats(findings)
        report = generate_threat_report(matches)
        j = format_threat_json(report)
        data = json.loads(j)
        assert "matches" in data
        assert "severity_counts" in data

    def test_format_threats_list(self):
        threats = _build_threat_db()[:5]
        text = format_threats_list_text(threats)
        assert "THREAT-001" in text
        assert "5 threats in database" in text

    def test_format_text_unmatched(self):
        report = ThreatReport(
            total_threats=5,
            unmatched_threats=["T-001", "T-002"],
        )
        text = format_threat_text(report)
        assert "Unmatched threats: 2" in text

    def test_format_text_severity_breakdown(self):
        threat = ThreatEntry(
            id="X", name="X", description="X",
            severity=ThreatSeverity.CRITICAL,
            mitigation="Fix",
        )
        report = ThreatReport(
            matches=[ThreatMatch(threat, "R", "T", "f.py", 1)],
            total_threats=1,
        )
        text = format_threat_text(report)
        assert "CRITICAL" in text
