"""Policy Engine for MCP Security Audit.

Evaluates MCP server configurations against customizable security
policies. Policies are defined as rule sets with conditions and actions.

Features:
- Policy definition with conditions and actions
- Policy sets with priority ordering
- Evaluation engine: ALLOW, DENY, WARN
- Built-in policies: OWASP Agentic, CIS-like, NIST-like
- Policy inheritance and composition
- Audit log of policy evaluations

CLI: mcpaudit policy --evaluate | --list | --audit

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

from __future__ import annotations

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


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

@unique
class PolicyAction(Enum):
    """Action to take when a policy condition matches."""
    ALLOW = "ALLOW"
    DENY = "DENY"
    WARN = "WARN"
    AUDIT = "AUDIT"  # log only


@unique
class PolicyResult(Enum):
    """Result of policy evaluation."""
    PASS = "PASS"
    FAIL = "FAIL"
    WARNING = "WARNING"
    NOT_APPLICABLE = "NOT_APPLICABLE"


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

@dataclass(frozen=True)
class PolicyCondition:
    """A condition to check in a policy rule."""
    field: str       # what to check (e.g., "tools", "sandbox_tier")
    operator: str    # how to check (e.g., "contains", "equals", "not_in")
    value: Any       # expected value
    description: str = ""

    def evaluate(self, context: dict[str, Any]) -> bool:
        """Evaluate condition against a context dict."""
        actual = context.get(self.field)
        if actual is None:
            return False

        if self.operator == "equals":
            return actual == self.value
        elif self.operator == "not_equals":
            return actual != self.value
        elif self.operator == "contains":
            if isinstance(actual, (list, tuple, set)):
                return self.value in actual
            return str(self.value) in str(actual)
        elif self.operator == "not_contains":
            if isinstance(actual, (list, tuple, set)):
                return self.value not in actual
            return str(self.value) not in str(actual)
        elif self.operator == "in":
            if isinstance(self.value, (list, tuple, set)):
                return actual in self.value
            return False
        elif self.operator == "not_in":
            if isinstance(self.value, (list, tuple, set)):
                return actual not in self.value
            return True
        elif self.operator == "greater_than":
            return float(actual) > float(self.value)
        elif self.operator == "less_than":
            return float(actual) < float(self.value)
        elif self.operator == "exists":
            return actual is not None
        return False

    def to_dict(self) -> dict[str, Any]:
        return {
            "field": self.field,
            "operator": self.operator,
            "value": self.value,
            "description": self.description,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PolicyCondition:
        return cls(
            field=data["field"],
            operator=data["operator"],
            value=data["value"],
            description=data.get("description", ""),
        )


@dataclass(frozen=True)
class PolicyRule:
    """A single policy rule with conditions and action."""
    rule_id: str
    name: str
    conditions: tuple[PolicyCondition, ...]
    action: PolicyAction
    severity: str = "MEDIUM"
    description: str = ""
    remediation: str = ""
    match_all: bool = True  # True=AND, False=OR

    def evaluate(self, context: dict[str, Any]) -> PolicyResult:
        """Evaluate this rule against a context."""
        if not self.conditions:
            return PolicyResult.NOT_APPLICABLE

        if self.match_all:
            matched = all(c.evaluate(context) for c in self.conditions)
        else:
            matched = any(c.evaluate(context) for c in self.conditions)

        if not matched:
            return PolicyResult.PASS

        if self.action == PolicyAction.DENY:
            return PolicyResult.FAIL
        elif self.action == PolicyAction.WARN:
            return PolicyResult.WARNING
        return PolicyResult.PASS

    def to_dict(self) -> dict[str, Any]:
        return {
            "rule_id": self.rule_id,
            "name": self.name,
            "conditions": [c.to_dict() for c in self.conditions],
            "action": self.action.value,
            "severity": self.severity,
            "description": self.description,
            "remediation": self.remediation,
            "match_all": self.match_all,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PolicyRule:
        return cls(
            rule_id=data["rule_id"],
            name=data["name"],
            conditions=tuple(
                PolicyCondition.from_dict(c) for c in data.get("conditions", [])
            ),
            action=PolicyAction(data.get("action", "DENY")),
            severity=data.get("severity", "MEDIUM"),
            description=data.get("description", ""),
            remediation=data.get("remediation", ""),
            match_all=data.get("match_all", True),
        )


@dataclass
class EvaluationResult:
    """Result of evaluating a rule against a context."""
    rule_id: str
    rule_name: str
    result: PolicyResult
    severity: str
    remediation: str = ""
    timestamp: str = ""

    def __post_init__(self) -> None:
        if not self.timestamp:
            self.timestamp = datetime.now(timezone.utc).isoformat()

    def to_dict(self) -> dict[str, Any]:
        return {
            "rule_id": self.rule_id,
            "rule_name": self.rule_name,
            "result": self.result.value,
            "severity": self.severity,
            "remediation": self.remediation,
            "timestamp": self.timestamp,
        }


# ---------------------------------------------------------------------------
# Engine
# ---------------------------------------------------------------------------

class PolicyEngine:
    """Evaluates configurations against security policies."""

    def __init__(self) -> None:
        self._rules: dict[str, PolicyRule] = {}
        self._audit_log: list[EvaluationResult] = []
        self._register_builtins()

    @property
    def rules(self) -> list[PolicyRule]:
        return list(self._rules.values())

    @property
    def rule_count(self) -> int:
        return len(self._rules)

    @property
    def audit_log(self) -> list[EvaluationResult]:
        return list(self._audit_log)

    def add_rule(self, rule: PolicyRule) -> None:
        self._rules[rule.rule_id] = rule

    def remove_rule(self, rule_id: str) -> bool:
        return self._rules.pop(rule_id, None) is not None

    def get_rule(self, rule_id: str) -> Optional[PolicyRule]:
        return self._rules.get(rule_id)

    def evaluate(self, context: dict[str, Any]) -> list[EvaluationResult]:
        """Evaluate all rules against a context."""
        results = []
        for rule in self._rules.values():
            result = rule.evaluate(context)
            eval_result = EvaluationResult(
                rule_id=rule.rule_id,
                rule_name=rule.name,
                result=result,
                severity=rule.severity,
                remediation=rule.remediation if result == PolicyResult.FAIL else "",
            )
            results.append(eval_result)
            self._audit_log.append(eval_result)
        return results

    def evaluate_batch(self, contexts: list[dict[str, Any]]) -> dict[str, list[EvaluationResult]]:
        """Evaluate rules against multiple contexts."""
        return {
            ctx.get("name", f"context_{i}"): self.evaluate(ctx)
            for i, ctx in enumerate(contexts)
        }

    def failures(self, results: list[EvaluationResult]) -> list[EvaluationResult]:
        """Filter to only failures."""
        return [r for r in results if r.result == PolicyResult.FAIL]

    def warnings(self, results: list[EvaluationResult]) -> list[EvaluationResult]:
        """Filter to only warnings."""
        return [r for r in results if r.result == PolicyResult.WARNING]

    def compliance_score(self, results: list[EvaluationResult]) -> float:
        """Compliance score: 1.0 = all pass, 0.0 = all fail."""
        applicable = [r for r in results if r.result != PolicyResult.NOT_APPLICABLE]
        if not applicable:
            return 1.0
        passed = sum(1 for r in applicable if r.result == PolicyResult.PASS)
        return passed / len(applicable)

    def clear_audit_log(self) -> None:
        self._audit_log.clear()

    # -- Persistence --

    def to_dict(self) -> dict[str, Any]:
        return {
            "version": "1.0",
            "rules": [r.to_dict() for r in self._rules.values()],
        }

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

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PolicyEngine:
        engine = cls()
        engine._rules.clear()
        for r in data.get("rules", []):
            rule = PolicyRule.from_dict(r)
            engine._rules[rule.rule_id] = rule
        return engine

    @classmethod
    def from_json(cls, json_str: str) -> PolicyEngine:
        return cls.from_dict(json.loads(json_str))

    # -- Report --

    def markdown_report(self, results: list[EvaluationResult]) -> str:
        """Generate policy evaluation report."""
        lines = ["# Policy Evaluation Report", ""]

        fails = self.failures(results)
        warns = self.warnings(results)
        score = self.compliance_score(results)

        lines.append("## Summary")
        lines.append(f"- **Rules evaluated**: {len(results)}")
        lines.append(f"- **Failures**: {len(fails)}")
        lines.append(f"- **Warnings**: {len(warns)}")
        lines.append(f"- **Compliance**: {score:.0%}")
        lines.append("")

        if fails:
            lines.append("## Failures")
            lines.append("| Rule | Severity | Remediation |")
            lines.append("|------|----------|-------------|")
            for r in fails:
                lines.append(f"| {r.rule_id}: {r.rule_name} | {r.severity} | {r.remediation or '-'} |")
            lines.append("")

        if warns:
            lines.append("## Warnings")
            for r in warns:
                lines.append(f"- **{r.rule_id}**: {r.rule_name}")
            lines.append("")

        return "\n".join(lines)

    # -- Built-in rules --

    def _register_builtins(self) -> None:
        """Register built-in security policy rules."""
        self.add_rule(PolicyRule(
            rule_id="POL-001",
            name="No FullAccess in production",
            conditions=(
                PolicyCondition("sandbox_tier", "equals", "FullAccess"),
                PolicyCondition("environment", "equals", "production"),
            ),
            action=PolicyAction.DENY,
            severity="CRITICAL",
            description="FullAccess sandbox tier is forbidden in production",
            remediation="Use WorkspaceWrite or ReadOnly in production",
        ))

        self.add_rule(PolicyRule(
            rule_id="POL-002",
            name="Bash requires constraints",
            conditions=(
                PolicyCondition("tools", "contains", "Bash"),
                PolicyCondition("constraint_count", "less_than", 1),
            ),
            action=PolicyAction.DENY,
            severity="HIGH",
            description="Agents with Bash must have at least one constraint",
            remediation="Add constraints to limit Bash execution scope",
        ))

        self.add_rule(PolicyRule(
            rule_id="POL-003",
            name="Web tools need review",
            conditions=(
                PolicyCondition("tools", "contains", "WebFetch"),
            ),
            action=PolicyAction.WARN,
            severity="MEDIUM",
            description="WebFetch can download arbitrary content",
            remediation="Consider URL allowlists",
        ))

        self.add_rule(PolicyRule(
            rule_id="POL-004",
            name="Excessive tool count",
            conditions=(
                PolicyCondition("tool_count", "greater_than", 8),
            ),
            action=PolicyAction.WARN,
            severity="LOW",
            description="Agent has many tools — consider splitting",
            remediation="Split into specialized agents with fewer tools",
        ))

        self.add_rule(PolicyRule(
            rule_id="POL-005",
            name="MCP without audit trail",
            conditions=(
                PolicyCondition("tools", "contains", "MCP"),
                PolicyCondition("has_audit", "equals", False),
            ),
            action=PolicyAction.DENY,
            severity="HIGH",
            description="MCP usage requires audit trail enabled",
            remediation="Enable audit trail for MCP interactions",
        ))
