"""Tests for MCP Vulnerability Database v2."""

from __future__ import annotations

import json
import tempfile
from pathlib import Path

import pytest

from mcpaudit.vuln_db_v2 import (
    EnrichedFinding,
    VulnDatabase,
    VulnEntry,
    VulnSearchResult,
    VulnSeverity,
    _build_vuln_db,
    _calculate_risk_score,
    export_db_json,
    export_db_to_file,
    format_enriched_json,
    format_enriched_text,
    format_vuln_db_json,
    format_vuln_db_text,
    import_db_from_file,
    import_db_json,
    match_findings,
)


# --- VulnSeverity ---

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

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

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

    def test_enum_iteration(self):
        assert len(list(VulnSeverity)) == 5


# --- VulnEntry ---

class TestVulnEntry:
    def test_creation(self):
        e = VulnEntry(id="V-001", title="Test", description="Desc",
                      severity=VulnSeverity.HIGH)
        assert e.id == "V-001"
        assert e.severity == VulnSeverity.HIGH

    def test_to_dict(self):
        e = VulnEntry(id="V-001", title="T", description="D",
                      severity=VulnSeverity.MEDIUM, tags=["a"], rule_ids=["R-001"],
                      cvss_score=5.0)
        d = e.to_dict()
        assert d["id"] == "V-001"
        assert d["severity"] == "MEDIUM"
        assert d["cvss_score"] == 5.0

    def test_from_dict(self):
        data = {"id": "V-002", "title": "T2", "description": "D2",
                "severity": "HIGH", "tags": ["x"], "cvss_score": 7.5}
        e = VulnEntry.from_dict(data)
        assert e.id == "V-002"
        assert e.severity == VulnSeverity.HIGH
        assert e.cvss_score == 7.5

    def test_from_dict_defaults(self):
        data = {"id": "V-003", "title": "T3", "description": "D3", "severity": "LOW"}
        e = VulnEntry.from_dict(data)
        assert e.tags == []
        assert e.cwe_ids == []
        assert e.exploitability == "UNKNOWN"
        assert e.cvss_score == 0.0

    def test_roundtrip(self):
        e = VulnEntry(id="V-R", title="RT", description="Roundtrip",
                      severity=VulnSeverity.CRITICAL, cwe_ids=["CWE-78"],
                      tags=["test"], published="2024-01-01", cvss_score=9.5)
        d = e.to_dict()
        e2 = VulnEntry.from_dict(d)
        assert e2.id == e.id
        assert e2.severity == e.severity
        assert e2.cvss_score == e.cvss_score

    def test_affected_patterns(self):
        e = VulnEntry(id="V-P", title="T", description="D",
                      severity=VulnSeverity.HIGH,
                      affected_patterns=["subprocess.call(", "os.system("])
        assert len(e.affected_patterns) == 2

    def test_affected_languages(self):
        e = VulnEntry(id="V-L", title="T", description="D",
                      severity=VulnSeverity.HIGH,
                      affected_languages=["python", "go"])
        assert "python" in e.affected_languages


# --- EnrichedFinding ---

class TestEnrichedFinding:
    def test_not_enriched(self):
        ef = EnrichedFinding(finding={"rule_id": "X"})
        assert not ef.is_enriched
        assert ef.max_severity == "INFO"

    def test_enriched_with_vulns(self):
        v = VulnEntry(id="V-1", title="T", description="D",
                      severity=VulnSeverity.CRITICAL)
        ef = EnrichedFinding(finding={"rule_id": "X"}, matched_vulns=[v])
        assert ef.is_enriched
        assert ef.max_severity == "CRITICAL"

    def test_max_severity_multiple(self):
        v1 = VulnEntry(id="V-1", title="T", description="D", severity=VulnSeverity.LOW)
        v2 = VulnEntry(id="V-2", title="T", description="D", severity=VulnSeverity.HIGH)
        ef = EnrichedFinding(finding={}, matched_vulns=[v1, v2])
        assert ef.max_severity == "HIGH"

    def test_to_dict(self):
        ef = EnrichedFinding(finding={"rule_id": "R"}, risk_score=42.0)
        d = ef.to_dict()
        assert d["risk_score"] == 42.0
        assert d["is_enriched"] is False

    def test_enrichment_notes(self):
        ef = EnrichedFinding(finding={}, enrichment_notes=["note1", "note2"])
        assert len(ef.enrichment_notes) == 2

    def test_finding_severity_fallback(self):
        ef = EnrichedFinding(finding={"severity": "HIGH"})
        assert ef.max_severity == "HIGH"


# --- VulnSearchResult ---

class TestVulnSearchResult:
    def test_empty(self):
        sr = VulnSearchResult()
        assert sr.count == 0

    def test_with_entries(self):
        e = VulnEntry(id="V", title="T", description="D", severity=VulnSeverity.LOW)
        sr = VulnSearchResult(entries=[e], query="test", total_in_db=10)
        assert sr.count == 1
        assert sr.total_in_db == 10

    def test_to_dict(self):
        sr = VulnSearchResult(query="q", total_in_db=5)
        d = sr.to_dict()
        assert d["query"] == "q"
        assert d["results_count"] == 0


# --- _build_vuln_db ---

class TestBuildVulnDB:
    def test_count(self):
        db = _build_vuln_db()
        assert len(db) >= 30

    def test_all_have_ids(self):
        for e in _build_vuln_db():
            assert e.id.startswith("MCP-VULN-")

    def test_all_have_severity(self):
        for e in _build_vuln_db():
            assert isinstance(e.severity, VulnSeverity)

    def test_all_have_cwe(self):
        for e in _build_vuln_db():
            assert len(e.cwe_ids) > 0

    def test_all_have_published(self):
        for e in _build_vuln_db():
            assert e.published != ""

    def test_unique_ids(self):
        ids = [e.id for e in _build_vuln_db()]
        assert len(ids) == len(set(ids))


# --- VulnDatabase ---

class TestVulnDatabase:
    def test_default_init(self):
        db = VulnDatabase()
        assert db.count >= 30

    def test_custom_entries(self):
        e = VulnEntry(id="V-1", title="T", description="D", severity=VulnSeverity.LOW)
        db = VulnDatabase(entries=[e])
        assert db.count == 1

    def test_get_existing(self):
        db = VulnDatabase()
        e = db.get("MCP-VULN-001")
        assert e is not None
        assert e.id == "MCP-VULN-001"

    def test_get_missing(self):
        db = VulnDatabase()
        assert db.get("NONEXISTENT") is None

    def test_list_all(self):
        db = VulnDatabase()
        all_entries = db.list_all()
        assert len(all_entries) == db.count

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

    def test_search_cwe(self):
        db = VulnDatabase()
        results = db.search("CWE-78")
        assert len(results) > 0

    def test_search_no_results(self):
        db = VulnDatabase()
        results = db.search("zzzznonexistentzzzz")
        assert len(results) == 0

    def test_search_case_insensitive(self):
        db = VulnDatabase()
        r1 = db.search("injection")
        r2 = db.search("INJECTION")
        assert len(r1) == len(r2)

    def test_filter_by_severity_enum(self):
        db = VulnDatabase()
        critical = db.filter_by_severity(VulnSeverity.CRITICAL)
        assert len(critical) > 0
        assert all(e.severity == VulnSeverity.CRITICAL for e in critical)

    def test_filter_by_severity_string(self):
        db = VulnDatabase()
        high = db.filter_by_severity("high")
        assert len(high) > 0

    def test_filter_by_cwe(self):
        db = VulnDatabase()
        results = db.filter_by_cwe("CWE-78")
        assert len(results) > 0
        assert all("CWE-78" in e.cwe_ids for e in results)

    def test_filter_by_cwe_no_prefix(self):
        db = VulnDatabase()
        results = db.filter_by_cwe("78")
        assert len(results) > 0

    def test_filter_by_tag(self):
        db = VulnDatabase()
        results = db.filter_by_tag("injection")
        assert len(results) > 0

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

    def test_filter_by_date_range(self):
        db = VulnDatabase()
        results = db.filter_by_date_range(date_from="2024-07-01", date_to="2024-09-30")
        assert len(results) > 0
        for e in results:
            assert "2024-07" <= e.published <= "2024-09-30"

    def test_filter_by_date_from_only(self):
        db = VulnDatabase()
        results = db.filter_by_date_range(date_from="2025-01-01")
        assert len(results) > 0

    def test_filter_by_date_to_only(self):
        db = VulnDatabase()
        results = db.filter_by_date_range(date_to="2024-06-30")
        assert len(results) > 0

    def test_vulns_for_rule(self):
        db = VulnDatabase()
        results = db.vulns_for_rule("CMD-001")
        assert len(results) > 0

    def test_vulns_for_rule_missing(self):
        db = VulnDatabase()
        results = db.vulns_for_rule("NONEXISTENT-999")
        assert results == []

    def test_add_entry(self):
        db = VulnDatabase(entries=[])
        e = VulnEntry(id="NEW-001", title="New", description="D",
                      severity=VulnSeverity.LOW, rule_ids=["R-NEW"])
        db.add_entry(e)
        assert db.count == 1
        assert db.get("NEW-001") is not None
        assert len(db.vulns_for_rule("R-NEW")) == 1

    def test_remove_entry(self):
        e = VulnEntry(id="DEL-001", title="Del", description="D",
                      severity=VulnSeverity.LOW, rule_ids=["R-DEL"])
        db = VulnDatabase(entries=[e])
        assert db.remove_entry("DEL-001") is True
        assert db.count == 0
        assert db.get("DEL-001") is None

    def test_remove_nonexistent(self):
        db = VulnDatabase(entries=[])
        assert db.remove_entry("NOPE") is False

    def test_search_advanced_query(self):
        db = VulnDatabase()
        r = db.search_advanced(query="injection")
        assert r.count > 0
        assert r.total_in_db == db.count

    def test_search_advanced_severity(self):
        db = VulnDatabase()
        r = db.search_advanced(severity="CRITICAL")
        assert r.count > 0

    def test_search_advanced_cwe(self):
        db = VulnDatabase()
        r = db.search_advanced(cwe="78")
        assert r.count > 0

    def test_search_advanced_language(self):
        db = VulnDatabase()
        r = db.search_advanced(language="python")
        assert r.count > 0

    def test_search_advanced_tag(self):
        db = VulnDatabase()
        r = db.search_advanced(tag="owasp")
        assert r.count > 0

    def test_search_advanced_date_range(self):
        db = VulnDatabase()
        r = db.search_advanced(date_from="2024-07-01", date_to="2024-12-31")
        assert r.count > 0

    def test_search_advanced_combined(self):
        db = VulnDatabase()
        r = db.search_advanced(severity="HIGH", language="python")
        assert r.count > 0

    def test_search_advanced_no_results(self):
        db = VulnDatabase()
        r = db.search_advanced(query="zzznonexistent", severity="INFO")
        assert r.count == 0


# --- match_findings ---

class TestMatchFindings:
    def test_match_by_rule_id(self):
        findings = [{"rule_id": "CMD-001", "title": "Command injection", "file": "a.py", "line": 1}]
        enriched = match_findings(findings)
        assert len(enriched) == 1
        assert enriched[0].is_enriched

    def test_no_match(self):
        findings = [{"rule_id": "NONEXISTENT-999", "title": "Nothing", "file": "b.py", "line": 1}]
        enriched = match_findings(findings)
        assert len(enriched) == 1
        assert not enriched[0].is_enriched

    def test_match_with_custom_db(self):
        e = VulnEntry(id="CUSTOM-1", title="Custom Vuln", description="D",
                      severity=VulnSeverity.HIGH, rule_ids=["CUSTOM-R"])
        db = VulnDatabase(entries=[e])
        findings = [{"rule_id": "CUSTOM-R", "title": "T", "file": "c.py", "line": 5}]
        enriched = match_findings(findings, db)
        assert enriched[0].is_enriched
        assert enriched[0].matched_vulns[0].id == "CUSTOM-1"

    def test_enrichment_notes(self):
        findings = [{"rule_id": "CMD-001", "title": "Cmd inj", "file": "a.py", "line": 1}]
        enriched = match_findings(findings)
        if enriched[0].is_enriched:
            assert len(enriched[0].enrichment_notes) > 0

    def test_risk_score_calculated(self):
        findings = [{"rule_id": "CMD-001", "title": "T", "file": "a.py", "line": 1}]
        enriched = match_findings(findings)
        assert enriched[0].risk_score > 0

    def test_multiple_findings(self):
        findings = [
            {"rule_id": "CMD-001", "title": "T1", "file": "a.py", "line": 1},
            {"rule_id": "NET-001", "title": "T2", "file": "b.py", "line": 2},
        ]
        enriched = match_findings(findings)
        assert len(enriched) == 2

    def test_empty_findings(self):
        enriched = match_findings([])
        assert enriched == []


# --- _calculate_risk_score ---

class TestCalculateRiskScore:
    def test_no_vulns_critical(self):
        s = _calculate_risk_score({"severity": "CRITICAL"}, [])
        assert s == 40.0

    def test_no_vulns_info(self):
        s = _calculate_risk_score({"severity": "INFO"}, [])
        assert s == 5.0

    def test_with_high_cvss(self):
        v = VulnEntry(id="V", title="T", description="D",
                      severity=VulnSeverity.CRITICAL, cvss_score=10.0,
                      exploitability="HIGH")
        s = _calculate_risk_score({}, [v])
        assert s >= 70.0

    def test_capped_at_100(self):
        vulns = [VulnEntry(id=f"V-{i}", title="T", description="D",
                           severity=VulnSeverity.CRITICAL, cvss_score=10.0,
                           exploitability="HIGH") for i in range(20)]
        s = _calculate_risk_score({}, vulns)
        assert s <= 100.0


# --- Import/Export ---

class TestImportExport:
    def test_export_json(self):
        db = VulnDatabase()
        j = export_db_json(db)
        data = json.loads(j)
        assert data["version"] == "2.0"
        assert data["count"] == db.count

    def test_import_json(self):
        db = VulnDatabase()
        j = export_db_json(db)
        db2 = import_db_json(j)
        assert db2.count == db.count

    def test_roundtrip_json(self):
        e = VulnEntry(id="RT-1", title="Roundtrip", description="D",
                      severity=VulnSeverity.HIGH, cvss_score=7.0)
        db = VulnDatabase(entries=[e])
        j = export_db_json(db)
        db2 = import_db_json(j)
        assert db2.get("RT-1") is not None
        assert db2.get("RT-1").cvss_score == 7.0

    def test_export_to_file(self):
        db = VulnDatabase(entries=[])
        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
            export_db_to_file(db, f.name)
            content = Path(f.name).read_text()
            assert "entries" in content

    def test_import_from_file(self):
        e = VulnEntry(id="FILE-1", title="File", description="D",
                      severity=VulnSeverity.LOW)
        db = VulnDatabase(entries=[e])
        with tempfile.NamedTemporaryFile(suffix=".json", delete=False, mode="w") as f:
            f.write(export_db_json(db))
            f.flush()
            db2 = import_db_from_file(f.name)
            assert db2.count == 1


# --- Formatting ---

class TestFormatting:
    def test_format_text(self):
        db = VulnDatabase()
        text = format_vuln_db_text(db)
        assert "Vulnerability Database" in text
        assert "MCP-VULN-001" in text

    def test_format_json(self):
        db = VulnDatabase()
        j = format_vuln_db_json(db)
        data = json.loads(j)
        assert "entries" in data

    def test_format_enriched_text(self):
        ef = EnrichedFinding(finding={"rule_id": "R", "title": "T"}, risk_score=50.0)
        text = format_enriched_text([ef])
        assert "Enriched" in text
        assert "50.0" in text

    def test_format_enriched_json(self):
        ef = EnrichedFinding(finding={"rule_id": "R"})
        j = format_enriched_json([ef])
        data = json.loads(j)
        assert len(data) == 1

    def test_format_enriched_empty(self):
        text = format_enriched_text([])
        assert "0" in text
