"""Tests for the Compliance Gap Analyzer module."""

from __future__ import annotations

import json

import pytest

from mcpaudit.compliance_gaps import (
    Control,
    ControlResult,
    ControlStatus,
    Framework,
    FrameworkResult,
    GapAnalysis,
    RemediationAction,
    RemediationPriority,
    Roadmap,
    analyze_gaps,
    format_gaps_html,
    format_gaps_json,
    format_gaps_text,
    generate_remediation_roadmap,
    get_framework_controls,
    list_supported_frameworks,
    _coverage_color,
    _escape_html,
    _evaluate_control,
    _extract_rule_prefixes,
)


# ---------------------------------------------------------------------------
# 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": "t1",
             "description": "d1", "file": "server.py", "line": 1},
            {"rule_id": "AUTH-001", "severity": "HIGH", "title": "t2",
             "description": "d2", "file": "server.py", "line": 2},
            {"rule_id": "RATE-001", "severity": "MEDIUM", "title": "t3",
             "description": "d3", "file": "tools.py", "line": 3},
            {"rule_id": "LOG-001", "severity": "LOW", "title": "t4",
             "description": "d4", "file": "tools.py", "line": 4},
        ]
    return {
        "summary": {
            "files_scanned": 2,
            "total_findings": total,
            "overall_grade": grade,
        },
        "results": [
            {"file": "server.py", "findings": findings[:2]},
            {"file": "tools.py", "findings": findings[2:]},
        ],
    }


def _rich_scan_data() -> dict:
    """Scan data with many different rule prefixes for good coverage."""
    findings = [
        {"rule_id": "CMD-001", "severity": "CRITICAL", "title": "t", "description": "d", "file": "f.py", "line": 1},
        {"rule_id": "AUTH-001", "severity": "HIGH", "title": "t", "description": "d", "file": "f.py", "line": 2},
        {"rule_id": "RATE-001", "severity": "MEDIUM", "title": "t", "description": "d", "file": "f.py", "line": 3},
        {"rule_id": "LOG-001", "severity": "LOW", "title": "t", "description": "d", "file": "f.py", "line": 4},
        {"rule_id": "PROMPT-001", "severity": "HIGH", "title": "t", "description": "d", "file": "f.py", "line": 5},
        {"rule_id": "INJ-001", "severity": "HIGH", "title": "t", "description": "d", "file": "f.py", "line": 6},
        {"rule_id": "TOOL-001", "severity": "MEDIUM", "title": "t", "description": "d", "file": "f.py", "line": 7},
        {"rule_id": "PRIV-001", "severity": "MEDIUM", "title": "t", "description": "d", "file": "f.py", "line": 8},
        {"rule_id": "LEAK-001", "severity": "HIGH", "title": "t", "description": "d", "file": "f.py", "line": 9},
        {"rule_id": "SEC-001", "severity": "LOW", "title": "t", "description": "d", "file": "f.py", "line": 10},
        {"rule_id": "CRYPTO-001", "severity": "MEDIUM", "title": "t", "description": "d", "file": "f.py", "line": 11},
        {"rule_id": "XSS-001", "severity": "MEDIUM", "title": "t", "description": "d", "file": "f.py", "line": 12},
        {"rule_id": "VALID-001", "severity": "LOW", "title": "t", "description": "d", "file": "f.py", "line": 13},
        {"rule_id": "SSRF-001", "severity": "HIGH", "title": "t", "description": "d", "file": "f.py", "line": 14},
        {"rule_id": "ERR-001", "severity": "LOW", "title": "t", "description": "d", "file": "f.py", "line": 15},
        {"rule_id": "SCOPE-001", "severity": "MEDIUM", "title": "t", "description": "d", "file": "f.py", "line": 16},
        {"rule_id": "MCP-001", "severity": "MEDIUM", "title": "t", "description": "d", "file": "f.py", "line": 17},
        {"rule_id": "PATH-001", "severity": "MEDIUM", "title": "t", "description": "d", "file": "f.py", "line": 18},
        {"rule_id": "DESER-001", "severity": "HIGH", "title": "t", "description": "d", "file": "f.py", "line": 19},
        {"rule_id": "CORS-001", "severity": "LOW", "title": "t", "description": "d", "file": "f.py", "line": 20},
        {"rule_id": "RESOURCE-001", "severity": "LOW", "title": "t", "description": "d", "file": "f.py", "line": 21},
        {"rule_id": "CONFUSED-001", "severity": "HIGH", "title": "t", "description": "d", "file": "f.py", "line": 22},
        {"rule_id": "AGENT-001", "severity": "HIGH", "title": "t", "description": "d", "file": "f.py", "line": 23},
    ]
    return {
        "summary": {"files_scanned": 1, "total_findings": len(findings), "overall_grade": "F"},
        "results": [{"file": "f.py", "findings": findings}],
    }


# ===========================================================================
# Framework
# ===========================================================================

class TestFramework:
    def test_all_values(self):
        assert len(Framework) == 6

    def test_from_string_owasp(self):
        assert Framework.from_string("owasp") == Framework.OWASP_AGENTIC
        assert Framework.from_string("owasp_agentic") == Framework.OWASP_AGENTIC
        assert Framework.from_string("owasp-agentic") == Framework.OWASP_AGENTIC

    def test_from_string_nist(self):
        assert Framework.from_string("nist") == Framework.NIST_AI_RMF
        assert Framework.from_string("nist_ai") == Framework.NIST_AI_RMF

    def test_from_string_eu(self):
        assert Framework.from_string("eu_ai_act") == Framework.EU_AI_ACT
        assert Framework.from_string("eu-ai-act") == Framework.EU_AI_ACT

    def test_from_string_soc2(self):
        assert Framework.from_string("soc2") == Framework.SOC2
        assert Framework.from_string("soc-2") == Framework.SOC2

    def test_from_string_hipaa(self):
        assert Framework.from_string("hipaa") == Framework.HIPAA

    def test_from_string_pci(self):
        assert Framework.from_string("pci_dss") == Framework.PCI_DSS
        assert Framework.from_string("pci-dss") == Framework.PCI_DSS
        assert Framework.from_string("pci") == Framework.PCI_DSS

    def test_from_string_case_insensitive(self):
        assert Framework.from_string("OWASP") == Framework.OWASP_AGENTIC

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

    def test_all_frameworks(self):
        fws = Framework.all_frameworks()
        assert len(fws) == 6
        assert Framework.OWASP_AGENTIC in fws


# ===========================================================================
# ControlStatus
# ===========================================================================

class TestControlStatus:
    def test_values(self):
        assert ControlStatus.COVERED.value == "Covered"
        assert ControlStatus.PARTIAL.value == "Partial"
        assert ControlStatus.GAP.value == "Gap"
        assert ControlStatus.NOT_APPLICABLE.value == "N/A"


# ===========================================================================
# RemediationPriority
# ===========================================================================

class TestRemediationPriority:
    def test_values(self):
        assert RemediationPriority.CRITICAL.value == "Critical"
        assert RemediationPriority.HIGH.value == "High"
        assert RemediationPriority.MEDIUM.value == "Medium"
        assert RemediationPriority.LOW.value == "Low"

    def test_order(self):
        assert RemediationPriority.CRITICAL.order < RemediationPriority.HIGH.order
        assert RemediationPriority.HIGH.order < RemediationPriority.MEDIUM.order
        assert RemediationPriority.MEDIUM.order < RemediationPriority.LOW.order


# ===========================================================================
# Control
# ===========================================================================

class TestControl:
    def test_to_dict(self):
        c = Control("C-1", "Auth", "desc", Framework.OWASP_AGENTIC, ["AUTH"], "Security")
        d = c.to_dict()
        assert d["control_id"] == "C-1"
        assert d["framework"] == "owasp"
        assert d["required_rules"] == ["AUTH"]

    def test_defaults(self):
        c = Control("C-1", "n", "d", Framework.SOC2)
        assert c.required_rules == []
        assert c.category == ""


# ===========================================================================
# ControlResult
# ===========================================================================

class TestControlResult:
    def test_to_dict(self):
        ctrl = Control("C-1", "Auth", "desc", Framework.SOC2, ["AUTH"])
        cr = ControlResult(ctrl, ControlStatus.COVERED, ["AUTH"], [], 100.0)
        d = cr.to_dict()
        assert d["status"] == "Covered"
        assert d["coverage_pct"] == 100.0

    def test_defaults(self):
        ctrl = Control("C-1", "n", "d", Framework.SOC2)
        cr = ControlResult(ctrl, ControlStatus.GAP)
        assert cr.matched_rules == []
        assert cr.missing_rules == []
        assert cr.coverage_pct == 0.0


# ===========================================================================
# FrameworkResult
# ===========================================================================

class TestFrameworkResult:
    def _make_result(self, statuses=None):
        if statuses is None:
            statuses = [ControlStatus.COVERED, ControlStatus.PARTIAL, ControlStatus.GAP]
        fr = FrameworkResult(framework=Framework.OWASP_AGENTIC)
        for i, status in enumerate(statuses):
            ctrl = Control(f"C-{i}", f"n{i}", "d", Framework.OWASP_AGENTIC, ["AUTH"])
            fr.controls.append(ControlResult(ctrl, status))
        return fr

    def test_counts(self):
        fr = self._make_result()
        assert fr.total_controls == 3
        assert fr.covered_controls == 1
        assert fr.partial_controls == 1
        assert fr.gap_controls == 1

    def test_coverage_pct(self):
        fr = self._make_result([ControlStatus.COVERED, ControlStatus.COVERED])
        assert fr.coverage_pct == 100.0

    def test_coverage_pct_partial(self):
        fr = self._make_result([ControlStatus.PARTIAL])
        assert fr.coverage_pct == 50.0

    def test_coverage_pct_gap(self):
        fr = self._make_result([ControlStatus.GAP])
        assert fr.coverage_pct == 0.0

    def test_coverage_pct_empty(self):
        fr = FrameworkResult(framework=Framework.SOC2)
        assert fr.coverage_pct == 0.0

    def test_coverage_pct_na_only(self):
        fr = self._make_result([ControlStatus.NOT_APPLICABLE])
        assert fr.coverage_pct == 100.0

    def test_to_dict(self):
        fr = self._make_result()
        d = fr.to_dict()
        assert d["framework"] == "owasp"
        assert d["total_controls"] == 3
        assert len(d["controls"]) == 3


# ===========================================================================
# Roadmap
# ===========================================================================

class TestRoadmap:
    def test_empty(self):
        r = Roadmap()
        assert len(r.actions) == 0
        assert r.total_estimated_days == 0

    def test_critical_actions(self):
        r = Roadmap()
        r.actions.append(RemediationAction(
            "R-1", "Fix auth", RemediationPriority.CRITICAL, Framework.SOC2
        ))
        r.actions.append(RemediationAction(
            "R-2", "Add logs", RemediationPriority.LOW, Framework.SOC2
        ))
        assert len(r.critical_actions) == 1

    def test_high_actions(self):
        r = Roadmap()
        r.actions.append(RemediationAction(
            "R-1", "Fix", RemediationPriority.HIGH, Framework.SOC2
        ))
        assert len(r.high_actions) == 1

    def test_to_dict(self):
        r = Roadmap()
        r.actions.append(RemediationAction(
            "R-1", "Fix", RemediationPriority.MEDIUM, Framework.SOC2,
            estimated_days=5,
        ))
        r.total_estimated_days = 5
        d = r.to_dict()
        assert d["total_actions"] == 1
        assert d["total_estimated_days"] == 5


# ===========================================================================
# GapAnalysis
# ===========================================================================

class TestGapAnalysis:
    def test_empty(self):
        ga = GapAnalysis()
        assert ga.overall_coverage == 0.0
        assert ga.total_gaps == 0

    def test_overall_coverage(self):
        ga = GapAnalysis()
        fr1 = FrameworkResult(Framework.OWASP_AGENTIC)
        ctrl = Control("C-1", "n", "d", Framework.OWASP_AGENTIC, ["AUTH"])
        fr1.controls.append(ControlResult(ctrl, ControlStatus.COVERED, coverage_pct=100))
        ga.frameworks.append(fr1)
        assert ga.overall_coverage > 0

    def test_total_gaps(self):
        ga = GapAnalysis()
        fr1 = FrameworkResult(Framework.OWASP_AGENTIC)
        ctrl = Control("C-1", "n", "d", Framework.OWASP_AGENTIC, ["AUTH"])
        fr1.controls.append(ControlResult(ctrl, ControlStatus.GAP))
        ga.frameworks.append(fr1)
        assert ga.total_gaps == 1

    def test_to_dict(self):
        ga = GapAnalysis(scan_grade="B", total_findings=3)
        d = ga.to_dict()
        assert d["scan_grade"] == "B"
        assert d["total_findings"] == 3
        assert "roadmap" in d


# ===========================================================================
# RemediationAction
# ===========================================================================

class TestRemediationAction:
    def test_to_dict(self):
        ra = RemediationAction(
            "R-1", "Fix auth", RemediationPriority.HIGH, Framework.OWASP_AGENTIC,
            ["C-1", "C-2"], "medium", 5,
        )
        d = ra.to_dict()
        assert d["action_id"] == "R-1"
        assert d["priority"] == "High"
        assert d["estimated_days"] == 5

    def test_defaults(self):
        ra = RemediationAction("R-1", "d", RemediationPriority.LOW, Framework.SOC2)
        assert ra.effort == "medium"
        assert ra.estimated_days == 5
        assert ra.control_ids == []


# ===========================================================================
# get_framework_controls
# ===========================================================================

class TestGetFrameworkControls:
    def test_owasp(self):
        controls = get_framework_controls(Framework.OWASP_AGENTIC)
        assert len(controls) >= 10
        assert all(c.framework == Framework.OWASP_AGENTIC for c in controls)

    def test_nist(self):
        controls = get_framework_controls(Framework.NIST_AI_RMF)
        assert len(controls) >= 5

    def test_eu_ai_act(self):
        controls = get_framework_controls(Framework.EU_AI_ACT)
        assert len(controls) >= 5

    def test_soc2(self):
        controls = get_framework_controls(Framework.SOC2)
        assert len(controls) >= 5

    def test_hipaa(self):
        controls = get_framework_controls(Framework.HIPAA)
        assert len(controls) >= 5

    def test_pci_dss(self):
        controls = get_framework_controls(Framework.PCI_DSS)
        assert len(controls) >= 5


class TestListSupportedFrameworks:
    def test_returns_all(self):
        fws = list_supported_frameworks()
        assert len(fws) == 6
        assert "owasp" in fws
        assert "nist" in fws


# ===========================================================================
# _extract_rule_prefixes
# ===========================================================================

class TestExtractRulePrefixes:
    def test_extracts_ids_and_prefixes(self):
        data = _scan_data()
        prefixes = _extract_rule_prefixes(data)
        assert "CMD-001" in prefixes
        assert "CMD" in prefixes
        assert "AUTH" in prefixes

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


# ===========================================================================
# _evaluate_control
# ===========================================================================

class TestEvaluateControl:
    def test_covered(self):
        ctrl = Control("C-1", "n", "d", Framework.SOC2, ["AUTH", "LOG"])
        cr = _evaluate_control(ctrl, {"AUTH-001", "AUTH", "LOG-001", "LOG"})
        assert cr.status == ControlStatus.COVERED
        assert cr.coverage_pct == 100.0

    def test_partial(self):
        ctrl = Control("C-1", "n", "d", Framework.SOC2, ["AUTH", "LOG"])
        cr = _evaluate_control(ctrl, {"AUTH-001", "AUTH"})
        assert cr.status == ControlStatus.PARTIAL
        assert 0 < cr.coverage_pct < 100

    def test_gap(self):
        ctrl = Control("C-1", "n", "d", Framework.SOC2, ["AUTH"])
        cr = _evaluate_control(ctrl, {"CMD-001", "CMD"})
        assert cr.status == ControlStatus.GAP
        assert cr.coverage_pct == 0.0

    def test_no_required_rules(self):
        ctrl = Control("C-1", "n", "d", Framework.SOC2, [])
        cr = _evaluate_control(ctrl, {"AUTH"})
        assert cr.status == ControlStatus.NOT_APPLICABLE


# ===========================================================================
# analyze_gaps
# ===========================================================================

class TestAnalyzeGaps:
    def test_single_framework(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        assert len(analysis.frameworks) == 1
        assert analysis.frameworks[0].framework == Framework.OWASP_AGENTIC

    def test_multiple_frameworks(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp", "nist", "soc2"])
        assert len(analysis.frameworks) == 3

    def test_all_frameworks(self):
        data = _scan_data()
        analysis = analyze_gaps(data, None)  # All frameworks
        assert len(analysis.frameworks) == 6

    def test_scan_grade_propagated(self):
        data = _scan_data(grade="A")
        analysis = analyze_gaps(data, ["owasp"])
        assert analysis.scan_grade == "A"

    def test_roadmap_generated(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        assert isinstance(analysis.roadmap, Roadmap)

    def test_gaps_detected(self):
        data = _scan_data(findings=[])  # No findings = no rule matches
        analysis = analyze_gaps(data, ["owasp"])
        assert analysis.total_gaps > 0

    def test_coverage_with_rich_scan(self):
        data = _rich_scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        assert analysis.overall_coverage > 0

    def test_empty_scan(self):
        data = {"summary": {}, "results": []}
        analysis = analyze_gaps(data, ["owasp"])
        assert analysis.scan_grade == "?"

    def test_returns_gap_analysis(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        assert isinstance(analysis, GapAnalysis)

    def test_to_dict(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        d = analysis.to_dict()
        assert "frameworks" in d
        assert "roadmap" in d
        assert "overall_coverage" in d


# ===========================================================================
# generate_remediation_roadmap
# ===========================================================================

class TestGenerateRemediationRoadmap:
    def test_generates_actions(self):
        data = _scan_data(findings=[])  # Gaps expected
        analysis = analyze_gaps(data, ["owasp"])
        roadmap = generate_remediation_roadmap(analysis)
        assert len(roadmap.actions) >= 1

    def test_actions_sorted_by_priority(self):
        data = _scan_data(findings=[])
        analysis = analyze_gaps(data, ["owasp"])
        roadmap = generate_remediation_roadmap(analysis)
        if len(roadmap.actions) >= 2:
            for i in range(len(roadmap.actions) - 1):
                assert roadmap.actions[i].priority.order <= roadmap.actions[i + 1].priority.order

    def test_total_days_calculated(self):
        data = _scan_data(findings=[])
        analysis = analyze_gaps(data, ["owasp"])
        roadmap = generate_remediation_roadmap(analysis)
        assert roadmap.total_estimated_days == sum(a.estimated_days for a in roadmap.actions)

    def test_no_gaps_no_actions(self):
        # When all controls are covered, no remediation needed
        analysis = GapAnalysis()
        fr = FrameworkResult(Framework.SOC2)
        ctrl = Control("C-1", "n", "d", Framework.SOC2, ["AUTH"])
        fr.controls.append(ControlResult(ctrl, ControlStatus.COVERED, ["AUTH"], []))
        analysis.frameworks.append(fr)
        roadmap = generate_remediation_roadmap(analysis)
        assert len(roadmap.actions) == 0


# ===========================================================================
# format_gaps_text
# ===========================================================================

class TestFormatGapsText:
    def test_basic_output(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        text = format_gaps_text(analysis)
        assert "Compliance Gap Analysis" in text
        assert "OWASP" in text

    def test_contains_coverage(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        text = format_gaps_text(analysis)
        assert "Coverage" in text
        assert "%" in text

    def test_contains_status_symbols(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        text = format_gaps_text(analysis)
        # Should have at least one status symbol
        assert any(sym in text for sym in ["[+]", "[~]", "[-]", "[.]"])

    def test_contains_roadmap(self):
        data = _scan_data(findings=[])
        analysis = analyze_gaps(data, ["owasp"])
        text = format_gaps_text(analysis)
        assert "Remediation Roadmap" in text


# ===========================================================================
# format_gaps_json
# ===========================================================================

class TestFormatGapsJson:
    def test_valid_json(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        j = format_gaps_json(analysis)
        parsed = json.loads(j)
        assert "frameworks" in parsed
        assert "roadmap" in parsed

    def test_round_trip(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        j = format_gaps_json(analysis)
        parsed = json.loads(j)
        assert parsed["scan_grade"] == analysis.scan_grade


# ===========================================================================
# format_gaps_html
# ===========================================================================

class TestFormatGapsHtml:
    def test_valid_html(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        html = format_gaps_html(analysis)
        assert "<!DOCTYPE html>" in html
        assert "</html>" in html

    def test_contains_title(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        html = format_gaps_html(analysis)
        assert "Compliance Gap Analysis" in html

    def test_contains_framework_cards(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp", "nist"])
        html = format_gaps_html(analysis)
        assert "fw-card" in html
        assert "OWASP" in html
        assert "NIST" in html

    def test_contains_checklist(self):
        data = _scan_data()
        analysis = analyze_gaps(data, ["owasp"])
        html = format_gaps_html(analysis)
        assert "Control Checklists" in html

    def test_contains_roadmap(self):
        data = _scan_data(findings=[])
        analysis = analyze_gaps(data, ["owasp"])
        html = format_gaps_html(analysis)
        assert "Remediation Roadmap" in html

    def test_coverage_color_high(self):
        assert _coverage_color(90.0) == "#22c55e"

    def test_coverage_color_medium(self):
        assert _coverage_color(65.0) == "#eab308"

    def test_coverage_color_low(self):
        assert _coverage_color(45.0) == "#f97316"

    def test_coverage_color_very_low(self):
        assert _coverage_color(20.0) == "#ef4444"


class TestEscapeHtmlGaps:
    def test_escapes(self):
        assert _escape_html("<b>") == "&lt;b&gt;"
        assert _escape_html("a & b") == "a &amp; b"
        assert _escape_html('"x"') == "&quot;x&quot;"
