"""Compliance Checker for MCP Security Audit.

Checks MCP server configurations against compliance frameworks
and generates compliance status reports.

Features:
- Framework definitions: OWASP Agentic, NIST AI RMF, EU AI Act, SOC2, ISO 27001
- Control mapping: which findings violate which controls
- Compliance score per framework (0-100)
- Gap analysis: which controls are not covered
- Evidence collection for audit trails
- Export as HTML, Markdown, JSON

CLI: mcpaudit compliance --framework owasp | --all | --gaps

Zero external dependencies. Part of mcpaudit v2.3.0.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from enum import Enum, unique
from typing import Any, Optional


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

@unique
class ControlStatus(Enum):
    """Status of a compliance control."""
    PASS = "PASS"
    FAIL = "FAIL"
    PARTIAL = "PARTIAL"
    NOT_TESTED = "NOT_TESTED"
    NOT_APPLICABLE = "NOT_APPLICABLE"


@unique
class Framework(Enum):
    """Supported compliance frameworks."""
    OWASP_AGENTIC = "owasp_agentic"
    NIST_AI_RMF = "nist_ai_rmf"
    EU_AI_ACT = "eu_ai_act"
    SOC2 = "soc2"
    ISO_27001 = "iso_27001"


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class Control:
    """A compliance control to check."""
    control_id: str
    name: str
    framework: Framework
    description: str = ""
    severity: str = "MEDIUM"
    related_rules: tuple[str, ...] = ()  # mcpaudit rule IDs

    def to_dict(self) -> dict[str, Any]:
        return {
            "control_id": self.control_id,
            "name": self.name,
            "framework": self.framework.value,
            "description": self.description,
            "severity": self.severity,
            "related_rules": list(self.related_rules),
        }


@dataclass
class ControlResult:
    """Result of checking a single control."""
    control: Control
    status: ControlStatus
    evidence: list[str] = field(default_factory=list)
    findings_count: int = 0

    def to_dict(self) -> dict[str, Any]:
        return {
            "control_id": self.control.control_id,
            "name": self.control.name,
            "status": self.status.value,
            "evidence": self.evidence,
            "findings_count": self.findings_count,
        }


@dataclass
class ComplianceReport:
    """Complete compliance report for a framework."""
    framework: Framework
    controls_checked: int = 0
    controls_passed: int = 0
    controls_failed: int = 0
    controls_partial: int = 0
    score: float = 0.0
    results: list[ControlResult] = field(default_factory=list)
    gaps: list[Control] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "framework": self.framework.value,
            "controls_checked": self.controls_checked,
            "controls_passed": self.controls_passed,
            "controls_failed": self.controls_failed,
            "controls_partial": self.controls_partial,
            "score": round(self.score, 1),
            "results": [r.to_dict() for r in self.results],
            "gaps": [g.to_dict() for g in self.gaps],
        }

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent)


# ---------------------------------------------------------------------------
# Built-in controls
# ---------------------------------------------------------------------------

_OWASP_CONTROLS = [
    Control("OA-01", "Excessive Agency Prevention", Framework.OWASP_AGENTIC,
            "Prevent agents from taking unauthorized actions",
            "CRITICAL", ("CMD-001", "CMD-002", "PRIV-001")),
    Control("OA-02", "Insecure Tool Implementation", Framework.OWASP_AGENTIC,
            "Ensure tools validate inputs and outputs",
            "CRITICAL", ("SEC-001", "SQL-001", "PATH-001")),
    Control("OA-03", "Insecure Output Handling", Framework.OWASP_AGENTIC,
            "Sanitize agent outputs before use",
            "HIGH", ("XSS-001", "DESER-001")),
    Control("OA-04", "Tool Misuse Detection", Framework.OWASP_AGENTIC,
            "Detect and prevent misuse of available tools",
            "HIGH", ("PRIV-001", "CMD-001")),
    Control("OA-05", "Insufficient Logging", Framework.OWASP_AGENTIC,
            "Log all agent actions for auditability",
            "MEDIUM", ("LOG-001",)),
]

_NIST_CONTROLS = [
    Control("NIST-G1", "Governance", Framework.NIST_AI_RMF,
            "Establish AI governance policies", "HIGH"),
    Control("NIST-M1", "Model Risk Management", Framework.NIST_AI_RMF,
            "Manage model selection and deployment risks", "HIGH"),
    Control("NIST-T1", "Transparency", Framework.NIST_AI_RMF,
            "Ensure transparency in AI operations", "MEDIUM"),
]

_EU_AI_CONTROLS = [
    Control("EU-HR1", "Human Oversight", Framework.EU_AI_ACT,
            "Ensure human oversight of AI systems", "CRITICAL"),
    Control("EU-T1", "Transparency", Framework.EU_AI_ACT,
            "Transparency requirements for AI systems", "HIGH"),
    Control("EU-R1", "Risk Assessment", Framework.EU_AI_ACT,
            "Conduct risk assessment for AI systems", "HIGH"),
]

_SOC2_CONTROLS = [
    Control("SOC2-CC1", "Control Environment", Framework.SOC2,
            "Establish and maintain control environment", "HIGH"),
    Control("SOC2-CC6", "Logical Access", Framework.SOC2,
            "Restrict logical access to information assets", "CRITICAL",
            ("PRIV-001", "PATH-001")),
    Control("SOC2-CC7", "System Operations", Framework.SOC2,
            "Monitor system operations", "HIGH", ("LOG-001",)),
]

_ISO_CONTROLS = [
    Control("ISO-A9", "Access Control", Framework.ISO_27001,
            "Information access control", "CRITICAL",
            ("PRIV-001", "PATH-001", "CMD-001")),
    Control("ISO-A12", "Operations Security", Framework.ISO_27001,
            "Operational security procedures", "HIGH"),
    Control("ISO-A18", "Compliance", Framework.ISO_27001,
            "Compliance with legal and contractual requirements", "HIGH"),
]

_ALL_CONTROLS: dict[Framework, list[Control]] = {
    Framework.OWASP_AGENTIC: _OWASP_CONTROLS,
    Framework.NIST_AI_RMF: _NIST_CONTROLS,
    Framework.EU_AI_ACT: _EU_AI_CONTROLS,
    Framework.SOC2: _SOC2_CONTROLS,
    Framework.ISO_27001: _ISO_CONTROLS,
}


# ---------------------------------------------------------------------------
# Checker
# ---------------------------------------------------------------------------

class ComplianceChecker:
    """Checks configurations against compliance frameworks."""

    def __init__(self) -> None:
        self._controls = dict(_ALL_CONTROLS)

    def get_frameworks(self) -> list[Framework]:
        return list(self._controls.keys())

    def get_controls(self, framework: Framework) -> list[Control]:
        return list(self._controls.get(framework, []))

    def add_control(self, control: Control) -> None:
        self._controls.setdefault(control.framework, []).append(control)

    def check_framework(self, framework: Framework,
                        findings: list[dict[str, Any]]) -> ComplianceReport:
        """Check all controls in a framework against scan findings."""
        controls = self._controls.get(framework, [])
        finding_rules = {f.get("rule_id", "") for f in findings}

        results = []
        passed = 0
        failed = 0
        partial = 0
        gaps = []

        for control in controls:
            if not control.related_rules:
                # No rules mapped — cannot test
                results.append(ControlResult(
                    control=control,
                    status=ControlStatus.NOT_TESTED,
                ))
                gaps.append(control)
                continue

            # Check if any related rules triggered findings
            triggered = [r for r in control.related_rules if r in finding_rules]
            matching_findings = [
                f for f in findings if f.get("rule_id", "") in control.related_rules
            ]

            if not triggered:
                results.append(ControlResult(
                    control=control,
                    status=ControlStatus.PASS,
                    evidence=[f"No violations found for rules: {', '.join(control.related_rules)}"],
                ))
                passed += 1
            elif len(triggered) == len(control.related_rules):
                results.append(ControlResult(
                    control=control,
                    status=ControlStatus.FAIL,
                    evidence=[f"Violations found: {', '.join(triggered)}"],
                    findings_count=len(matching_findings),
                ))
                failed += 1
            else:
                results.append(ControlResult(
                    control=control,
                    status=ControlStatus.PARTIAL,
                    evidence=[f"Some violations: {', '.join(triggered)}"],
                    findings_count=len(matching_findings),
                ))
                partial += 1

        checked = passed + failed + partial
        score = (passed / checked * 100) if checked > 0 else 0.0

        return ComplianceReport(
            framework=framework,
            controls_checked=checked,
            controls_passed=passed,
            controls_failed=failed,
            controls_partial=partial,
            score=score,
            results=results,
            gaps=gaps,
        )

    def check_all(self, findings: list[dict[str, Any]]) -> list[ComplianceReport]:
        """Check all frameworks."""
        return [
            self.check_framework(fw, findings)
            for fw in self._controls
        ]

    def markdown_report(self, report: ComplianceReport) -> str:
        """Generate compliance report."""
        lines = [f"# Compliance Report: {report.framework.value.upper()}", ""]

        lines.append("## Summary")
        lines.append(f"- **Score**: {report.score:.0f}%")
        lines.append(f"- **Controls checked**: {report.controls_checked}")
        lines.append(f"- **Passed**: {report.controls_passed}")
        lines.append(f"- **Failed**: {report.controls_failed}")
        lines.append(f"- **Partial**: {report.controls_partial}")
        lines.append("")

        if report.results:
            lines.append("## Controls")
            lines.append("| ID | Name | Status | Findings |")
            lines.append("|----|------|--------|----------|")
            for r in report.results:
                lines.append(
                    f"| {r.control.control_id} | {r.control.name} "
                    f"| {r.status.value} | {r.findings_count} |"
                )
            lines.append("")

        if report.gaps:
            lines.append("## Coverage Gaps")
            for g in report.gaps:
                lines.append(f"- **{g.control_id}**: {g.name} — no rules mapped")
            lines.append("")

        return "\n".join(lines)
