"""MCP Compliance Validator.

Validates MCP server scan results against security compliance frameworks:
OWASP Agentic Top 10, NIST AI RMF, and SOC2 AI controls.
Maps mcpaudit rules to framework controls and provides pass/fail reports.
Zero external dependencies.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any


class ComplianceFramework(Enum):
    """Supported compliance frameworks."""

    OWASP_AGENTIC = "owasp"
    NIST_AI = "nist"
    SOC2_AI = "soc2"

    @classmethod
    def from_string(cls, value: str) -> ComplianceFramework:
        """Parse a framework string (case-insensitive)."""
        mapping = {
            "owasp": cls.OWASP_AGENTIC,
            "owasp_agentic": cls.OWASP_AGENTIC,
            "owasp-agentic": cls.OWASP_AGENTIC,
            "nist": cls.NIST_AI,
            "nist_ai": cls.NIST_AI,
            "nist-ai": cls.NIST_AI,
            "soc2": cls.SOC2_AI,
            "soc2_ai": cls.SOC2_AI,
            "soc2-ai": cls.SOC2_AI,
        }
        key = value.lower().strip()
        if key in mapping:
            return mapping[key]
        valid = ", ".join(sorted(mapping.keys()))
        raise ValueError(f"Unknown framework '{value}'. Valid: {valid}")


class ControlStatus(Enum):
    """Status of a compliance control check."""

    PASS = "PASS"
    FAIL = "FAIL"
    NOT_APPLICABLE = "N/A"
    PARTIAL = "PARTIAL"


@dataclass
class ComplianceCheck:
    """A single compliance control check."""

    control_id: str
    name: str
    description: str
    how_to_validate: str
    scanner_rules: list[str] = field(default_factory=list)
    severity: str = "HIGH"

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "control_id": self.control_id,
            "name": self.name,
            "description": self.description,
            "how_to_validate": self.how_to_validate,
            "scanner_rules": self.scanner_rules,
            "severity": self.severity,
        }


@dataclass
class ControlResult:
    """Result of evaluating a single control."""

    check: ComplianceCheck
    status: ControlStatus
    findings_count: int = 0
    finding_details: list[dict[str, Any]] = field(default_factory=list)
    notes: str = ""

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "control_id": self.check.control_id,
            "name": self.check.name,
            "status": self.status.value,
            "findings_count": self.findings_count,
            "finding_details": self.finding_details,
            "notes": self.notes,
        }


@dataclass
class ComplianceResult:
    """Full compliance validation result."""

    framework: ComplianceFramework
    controls: list[ControlResult] = field(default_factory=list)
    scan_file: str = ""
    timestamp: str = ""

    @property
    def total_controls(self) -> int:
        return len(self.controls)

    @property
    def passed(self) -> int:
        return sum(1 for c in self.controls if c.status == ControlStatus.PASS)

    @property
    def failed(self) -> int:
        return sum(1 for c in self.controls if c.status == ControlStatus.FAIL)

    @property
    def partial(self) -> int:
        return sum(1 for c in self.controls if c.status == ControlStatus.PARTIAL)

    @property
    def not_applicable(self) -> int:
        return sum(
            1 for c in self.controls if c.status == ControlStatus.NOT_APPLICABLE
        )

    @property
    def compliance_score(self) -> float:
        """Percentage of applicable controls that pass."""
        applicable = self.total_controls - self.not_applicable
        if applicable == 0:
            return 100.0
        passing = self.passed + (self.partial * 0.5)
        return round((passing / applicable) * 100, 1)

    @property
    def overall_status(self) -> str:
        """Overall compliance verdict."""
        if self.failed == 0:
            return "COMPLIANT"
        if self.compliance_score >= 80:
            return "PARTIALLY_COMPLIANT"
        return "NON_COMPLIANT"

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "framework": self.framework.value,
            "scan_file": self.scan_file,
            "overall_status": self.overall_status,
            "compliance_score": self.compliance_score,
            "summary": {
                "total": self.total_controls,
                "passed": self.passed,
                "failed": self.failed,
                "partial": self.partial,
                "not_applicable": self.not_applicable,
            },
            "controls": [c.to_dict() for c in self.controls],
        }


# ---------------------------------------------------------------------------
# Framework Definitions
# ---------------------------------------------------------------------------

def _owasp_agentic_checks() -> list[ComplianceCheck]:
    """OWASP Agentic Top 10 controls."""
    return [
        ComplianceCheck(
            control_id="OWASP-A01",
            name="Confused Deputy Prevention",
            description=(
                "Ensure tools verify authorization and do not act solely "
                "on LLM-provided context for privileged operations."
            ),
            how_to_validate="Check for authorization controls in tool handlers",
            scanner_rules=["AUTH-001", "AUTH-002", "AUTH-003", "MCP-AUTH-001"],
            severity="CRITICAL",
        ),
        ComplianceCheck(
            control_id="OWASP-A02",
            name="Prompt Injection Defense",
            description=(
                "Protect against prompt injection through tool inputs "
                "and outputs that may manipulate LLM behavior."
            ),
            how_to_validate="Verify input/output sanitization in tool handlers",
            scanner_rules=[
                "INJ-001", "INJ-002", "INJ-003", "MCP-INJ-001", "MCP-INJ-002",
            ],
            severity="CRITICAL",
        ),
        ComplianceCheck(
            control_id="OWASP-A03",
            name="Tool Integrity Verification",
            description=(
                "Verify integrity of tools and prevent poisoning through "
                "tampered tool definitions or implementations."
            ),
            how_to_validate="Check for tool pinning, checksums, and integrity checks",
            scanner_rules=["TOOL-001", "TOOL-002", "MCP-TOOL-001"],
            severity="CRITICAL",
        ),
        ComplianceCheck(
            control_id="OWASP-A04",
            name="Least Privilege Enforcement",
            description=(
                "Ensure tools have minimal permissions necessary "
                "and do not over-request access."
            ),
            how_to_validate="Verify permission scoping in tool configurations",
            scanner_rules=["AUTH-001", "PERM-001"],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="OWASP-A05",
            name="Rug Pull Prevention",
            description=(
                "Prevent runtime modification of tool definitions "
                "after initial trust establishment."
            ),
            how_to_validate="Check for immutable tool definitions and change detection",
            scanner_rules=["TOOL-001", "TOOL-002", "MCP-TOOL-001"],
            severity="CRITICAL",
        ),
        ComplianceCheck(
            control_id="OWASP-A06",
            name="Data Exfiltration Prevention",
            description=(
                "Prevent unauthorized data transmission through "
                "tool outputs, network calls, or side channels."
            ),
            how_to_validate="Verify egress controls and output sanitization",
            scanner_rules=[
                "NET-001", "NET-002", "EXFIL-001", "MCP-NET-001",
            ],
            severity="CRITICAL",
        ),
        ComplianceCheck(
            control_id="OWASP-A07",
            name="Tool Chain Security",
            description=(
                "Prevent privilege escalation through tool chaining "
                "and verify composite permissions."
            ),
            how_to_validate=(
                "Check for workflow-level authorization and cumulative action tracking"
            ),
            scanner_rules=["AUTH-001", "AUTH-002", "CHAIN-001"],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="OWASP-A08",
            name="Input Validation",
            description=(
                "Validate and sanitize all tool inputs to prevent "
                "injection attacks and unexpected behavior."
            ),
            how_to_validate="Check for input validation in tool handlers",
            scanner_rules=["CMD-001", "CMD-002", "CMD-003", "SQL-001", "SQL-002"],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="OWASP-A09",
            name="Security Logging and Monitoring",
            description=(
                "Log all tool invocations, security events, and errors "
                "for audit trails and incident detection."
            ),
            how_to_validate="Verify logging in tool handlers and security events",
            scanner_rules=["LOG-001", "LOG-002", "MCP-LOG-001"],
            severity="MEDIUM",
        ),
        ComplianceCheck(
            control_id="OWASP-A10",
            name="Rate Limiting and Resource Control",
            description=(
                "Implement rate limiting and resource bounds to prevent "
                "denial of service and resource exhaustion."
            ),
            how_to_validate="Check for rate limiting and resource controls",
            scanner_rules=["RATE-001", "RATE-002", "DOS-001"],
            severity="MEDIUM",
        ),
    ]


def _nist_ai_checks() -> list[ComplianceCheck]:
    """NIST AI RMF relevant controls for MCP."""
    return [
        ComplianceCheck(
            control_id="NIST-GOV-1",
            name="AI System Governance",
            description=(
                "Establish governance for AI tool usage including access "
                "controls, policies, and oversight mechanisms."
            ),
            how_to_validate="Verify authorization and policy enforcement",
            scanner_rules=["AUTH-001", "AUTH-002", "PERM-001"],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="NIST-MAP-1",
            name="Risk Identification",
            description=(
                "Identify and document risks in AI tool interactions "
                "including injection, data leaks, and misuse."
            ),
            how_to_validate="Check for vulnerability scanning coverage",
            scanner_rules=[
                "INJ-001", "INJ-002", "CMD-001", "NET-001", "SQL-001",
            ],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="NIST-MEA-1",
            name="Security Measurement",
            description=(
                "Measure and monitor security posture of AI tools "
                "through continuous scanning and metrics."
            ),
            how_to_validate="Verify logging and monitoring capabilities",
            scanner_rules=["LOG-001", "LOG-002", "MCP-LOG-001"],
            severity="MEDIUM",
        ),
        ComplianceCheck(
            control_id="NIST-MAN-1",
            name="Risk Management",
            description=(
                "Manage identified risks through remediation, mitigation, "
                "and continuous improvement."
            ),
            how_to_validate="Check for remediation and fix capabilities",
            scanner_rules=["CMD-001", "CMD-002", "PATH-001", "SER-001"],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="NIST-SEC-1",
            name="Secure Development",
            description=(
                "Follow secure development practices for AI tools "
                "including input validation and output sanitization."
            ),
            how_to_validate="Verify secure coding patterns",
            scanner_rules=[
                "CMD-001", "CMD-002", "CMD-003", "INJ-001", "SQL-001",
            ],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="NIST-PRI-1",
            name="Data Privacy",
            description=(
                "Protect sensitive data in AI tool interactions, prevent "
                "unauthorized data access and exfiltration."
            ),
            how_to_validate="Check for data protection controls",
            scanner_rules=[
                "SEC-001", "SEC-002", "CRED-001", "NET-001", "EXFIL-001",
            ],
            severity="CRITICAL",
        ),
    ]


def _soc2_ai_checks() -> list[ComplianceCheck]:
    """SOC2 Type II AI-relevant controls for MCP."""
    return [
        ComplianceCheck(
            control_id="SOC2-CC6.1",
            name="Logical Access Controls",
            description=(
                "Implement logical access controls for AI tools to "
                "prevent unauthorized access and privilege escalation."
            ),
            how_to_validate="Verify authentication and authorization in tools",
            scanner_rules=["AUTH-001", "AUTH-002", "AUTH-003", "PERM-001"],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="SOC2-CC6.3",
            name="Access Restriction",
            description=(
                "Restrict access to AI tools based on roles and "
                "enforce least privilege principles."
            ),
            how_to_validate="Check for role-based access controls",
            scanner_rules=["AUTH-001", "PERM-001", "MCP-AUTH-001"],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="SOC2-CC7.1",
            name="Security Monitoring",
            description=(
                "Monitor AI tool operations for security anomalies, "
                "unauthorized access, and data breaches."
            ),
            how_to_validate="Verify security event logging and monitoring",
            scanner_rules=["LOG-001", "LOG-002", "MCP-LOG-001"],
            severity="MEDIUM",
        ),
        ComplianceCheck(
            control_id="SOC2-CC7.2",
            name="Incident Detection",
            description=(
                "Detect and respond to security incidents in AI tool "
                "operations including injection and exfiltration."
            ),
            how_to_validate="Check for incident detection capabilities",
            scanner_rules=[
                "INJ-001", "INJ-002", "NET-001", "EXFIL-001",
            ],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="SOC2-CC8.1",
            name="Change Management",
            description=(
                "Control changes to AI tools through version pinning, "
                "integrity verification, and change logging."
            ),
            how_to_validate="Verify tool integrity and change controls",
            scanner_rules=["TOOL-001", "TOOL-002", "MCP-TOOL-001"],
            severity="HIGH",
        ),
        ComplianceCheck(
            control_id="SOC2-C1.1",
            name="Data Confidentiality",
            description=(
                "Protect confidentiality of data processed by AI tools, "
                "prevent leaks through secrets, network, and storage."
            ),
            how_to_validate="Check for data protection and encryption",
            scanner_rules=[
                "SEC-001", "SEC-002", "CRED-001", "TLS-001", "CRYPTO-001",
            ],
            severity="CRITICAL",
        ),
    ]


def get_framework_checks(framework: ComplianceFramework) -> list[ComplianceCheck]:
    """Get compliance checks for a given framework."""
    builders = {
        ComplianceFramework.OWASP_AGENTIC: _owasp_agentic_checks,
        ComplianceFramework.NIST_AI: _nist_ai_checks,
        ComplianceFramework.SOC2_AI: _soc2_ai_checks,
    }
    builder = builders.get(framework)
    if builder is None:
        raise ValueError(f"Unknown framework: {framework}")
    return builder()


def _extract_findings(scan_data: dict[str, Any]) -> list[dict[str, Any]]:
    """Extract flat list of findings from scan data."""
    findings: list[dict[str, Any]] = []
    results = scan_data.get("results", [])
    for result in results:
        for f in result.get("findings", []):
            findings.append(f)
    return findings


def _evaluate_control(
    check: ComplianceCheck,
    findings: list[dict[str, Any]],
) -> ControlResult:
    """Evaluate a single compliance control against findings.

    A control PASSES if no findings match its scanner_rules.
    A control FAILS if any findings match.
    A control is PARTIAL if some rules have findings but others don't.
    """
    matched_findings: list[dict[str, Any]] = []
    matched_rules: set[str] = set()

    for finding in findings:
        rule_id = finding.get("rule_id", "")
        if rule_id in check.scanner_rules:
            matched_findings.append(finding)
            matched_rules.add(rule_id)

    if not matched_findings:
        return ControlResult(
            check=check,
            status=ControlStatus.PASS,
            findings_count=0,
            notes="No violations detected",
        )

    # If only some rules triggered, it's partial
    if len(matched_rules) < len(check.scanner_rules) and len(check.scanner_rules) > 1:
        status = ControlStatus.PARTIAL
        notes = (
            f"{len(matched_rules)}/{len(check.scanner_rules)} "
            f"rules triggered with {len(matched_findings)} finding(s)"
        )
    else:
        status = ControlStatus.FAIL
        notes = f"{len(matched_findings)} violation(s) detected"

    return ControlResult(
        check=check,
        status=status,
        findings_count=len(matched_findings),
        finding_details=[
            {
                "rule_id": f.get("rule_id", ""),
                "title": f.get("title", ""),
                "file": f.get("file", ""),
                "line": f.get("line", 0),
                "severity": f.get("severity", "UNKNOWN"),
            }
            for f in matched_findings
        ],
        notes=notes,
    )


def validate_compliance(
    scan_data: dict[str, Any],
    framework: ComplianceFramework,
) -> ComplianceResult:
    """Validate scan results against a compliance framework.

    Args:
        scan_data: Scan result dict with 'results' list.
        framework: Compliance framework to validate against.

    Returns:
        ComplianceResult with per-control pass/fail status.
    """
    checks = get_framework_checks(framework)
    findings = _extract_findings(scan_data)
    controls = [_evaluate_control(check, findings) for check in checks]

    return ComplianceResult(
        framework=framework,
        controls=controls,
    )


def validate_compliance_from_file(
    scan_file: Path,
    framework: ComplianceFramework,
) -> ComplianceResult:
    """Validate compliance from a scan result JSON file.

    Args:
        scan_file: Path to scan result JSON file.
        framework: Compliance framework to validate against.

    Returns:
        ComplianceResult with per-control pass/fail status.
    """
    content = scan_file.read_text(encoding="utf-8")
    scan_data = json.loads(content)
    result = validate_compliance(scan_data, framework)
    result.scan_file = str(scan_file)
    return result


def format_compliance_text(result: ComplianceResult) -> str:
    """Format compliance result as human-readable text."""
    lines: list[str] = []
    lines.append("=" * 65)
    lines.append(f"MCP Compliance Report — {result.framework.value.upper()}")
    lines.append("=" * 65)
    lines.append(f"Overall: {result.overall_status}")
    lines.append(f"Score:   {result.compliance_score}%")
    lines.append(
        f"Controls: {result.passed} passed, {result.failed} failed, "
        f"{result.partial} partial, {result.not_applicable} N/A"
    )
    lines.append("")

    status_icon = {
        ControlStatus.PASS: "[PASS]",
        ControlStatus.FAIL: "[FAIL]",
        ControlStatus.PARTIAL: "[PART]",
        ControlStatus.NOT_APPLICABLE: "[ NA ]",
    }

    for c in result.controls:
        icon = status_icon.get(c.status, "[????]")
        lines.append(f"  {icon} {c.check.control_id}: {c.check.name}")
        if c.findings_count > 0:
            lines.append(f"         {c.notes}")
            for fd in c.finding_details[:3]:
                lines.append(
                    f"         - {fd.get('rule_id', '')}: "
                    f"{fd.get('title', '')} ({fd.get('file', '')}:"
                    f"{fd.get('line', 0)})"
                )
            if len(c.finding_details) > 3:
                lines.append(
                    f"         ... and {len(c.finding_details) - 3} more"
                )
    lines.append("")
    return "\n".join(lines)


def format_compliance_json(result: ComplianceResult) -> str:
    """Format compliance result as JSON."""
    return json.dumps(result.to_dict(), indent=2)


def format_compliance_html(result: ComplianceResult) -> str:
    """Format compliance result as self-contained HTML."""
    rows = ""
    for c in result.controls:
        color = {
            ControlStatus.PASS: "#28a745",
            ControlStatus.FAIL: "#dc3545",
            ControlStatus.PARTIAL: "#ffc107",
            ControlStatus.NOT_APPLICABLE: "#6c757d",
        }.get(c.status, "#333")

        details = ""
        if c.finding_details:
            items = "".join(
                f"<li>{fd.get('rule_id', '')}: {fd.get('title', '')}</li>"
                for fd in c.finding_details[:5]
            )
            details = f"<ul>{items}</ul>"

        rows += (
            f"<tr>"
            f"<td>{c.check.control_id}</td>"
            f"<td>{c.check.name}</td>"
            f'<td style="color:{color};font-weight:bold">{c.status.value}</td>'
            f"<td>{c.findings_count}</td>"
            f"<td>{details}</td>"
            f"</tr>\n"
        )

    score_color = "#28a745" if result.compliance_score >= 80 else "#dc3545"

    return f"""<!DOCTYPE html>
<html>
<head><title>MCP Compliance Report</title>
<style>
body {{ font-family: sans-serif; margin: 2em; }}
table {{ border-collapse: collapse; width: 100%; }}
th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
th {{ background: #333; color: #fff; }}
tr:nth-child(even) {{ background: #f9f9f9; }}
.score {{ font-size: 2em; color: {score_color}; }}
</style>
</head>
<body>
<h1>MCP Compliance Report: {result.framework.value.upper()}</h1>
<p class="score">{result.compliance_score}% — {result.overall_status}</p>
<p>Passed: {result.passed} | Failed: {result.failed} | \
Partial: {result.partial} | N/A: {result.not_applicable}</p>
<table>
<tr><th>Control</th><th>Name</th><th>Status</th>\
<th>Findings</th><th>Details</th></tr>
{rows}
</table>
</body>
</html>"""


def list_frameworks() -> list[dict[str, Any]]:
    """List available compliance frameworks with control counts."""
    result = []
    for fw in ComplianceFramework:
        checks = get_framework_checks(fw)
        result.append({
            "id": fw.value,
            "name": fw.name,
            "controls": len(checks),
        })
    return result
