"""Security Policy Engine v2 — Advanced organizational policy evaluation.

Evaluates security findings against hierarchical, YAML-like policy rules
(parsed as JSON dicts) with rich operators and composable actions.

Features:
- Policy language: YAML-like rules parsed as dict/JSON
- Operators: eq, neq, gt, lt, gte, lte, in, not_in, contains, matches
- Actions: BLOCK (fail CI), WARN (advisory), ALLOW (whitelist), REQUIRE (mandatory)
- Policy inheritance: org -> team -> project
- Merge semantics: child overrides parent by rule_id, inherits the rest
- Batch evaluation across multiple finding sets

CLI: mcpaudit policy-v2 --evaluate | --merge | --validate

Zero external dependencies. Python 3.11+ stdlib only.
Part of mcpaudit v3.0.0.
"""

from __future__ import annotations

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


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

@unique
class PolicyV2Action(Enum):
    """Action to take when a policy rule matches."""
    BLOCK = "BLOCK"
    WARN = "WARN"
    ALLOW = "ALLOW"
    REQUIRE = "REQUIRE"


@unique
class PolicyV2Verdict(Enum):
    """Verdict for a single finding against a policy."""
    BLOCKED = "BLOCKED"
    WARNED = "WARNED"
    ALLOWED = "ALLOWED"
    REQUIRED_MISSING = "REQUIRED_MISSING"
    PASSED = "PASSED"
    NOT_EVALUATED = "NOT_EVALUATED"


# ---------------------------------------------------------------------------
# Operators
# ---------------------------------------------------------------------------

_OPERATORS: dict[str, Any] = {}


def _register_op(name: str):
    """Decorator to register an operator function."""
    def decorator(fn):
        _OPERATORS[name] = fn
        return fn
    return decorator


@_register_op("eq")
def _op_eq(actual: Any, expected: Any) -> bool:
    if isinstance(actual, str) and isinstance(expected, str):
        return actual.lower() == expected.lower()
    return actual == expected


@_register_op("neq")
def _op_neq(actual: Any, expected: Any) -> bool:
    return not _op_eq(actual, expected)


@_register_op("gt")
def _op_gt(actual: Any, expected: Any) -> bool:
    try:
        return float(actual) > float(expected)
    except (ValueError, TypeError):
        return False


@_register_op("lt")
def _op_lt(actual: Any, expected: Any) -> bool:
    try:
        return float(actual) < float(expected)
    except (ValueError, TypeError):
        return False


@_register_op("gte")
def _op_gte(actual: Any, expected: Any) -> bool:
    try:
        return float(actual) >= float(expected)
    except (ValueError, TypeError):
        return False


@_register_op("lte")
def _op_lte(actual: Any, expected: Any) -> bool:
    try:
        return float(actual) <= float(expected)
    except (ValueError, TypeError):
        return False


@_register_op("in")
def _op_in(actual: Any, expected: Any) -> bool:
    if isinstance(expected, (list, tuple, set)):
        return actual in expected
    return False


@_register_op("not_in")
def _op_not_in(actual: Any, expected: Any) -> bool:
    if isinstance(expected, (list, tuple, set)):
        return actual not in expected
    return True


@_register_op("contains")
def _op_contains(actual: Any, expected: Any) -> bool:
    if isinstance(actual, (list, tuple, set)):
        return expected in actual
    if isinstance(actual, str):
        return str(expected) in actual
    return False


@_register_op("matches")
def _op_matches(actual: Any, expected: Any) -> bool:
    if not isinstance(actual, str) or not isinstance(expected, str):
        return False
    try:
        return bool(re.search(expected, actual))
    except re.error:
        return False


def evaluate_operator(op: str, actual: Any, expected: Any) -> bool:
    """Evaluate a named operator against actual and expected values."""
    fn = _OPERATORS.get(op)
    if fn is None:
        return False
    return fn(actual, expected)


def get_supported_operators() -> list[str]:
    """Return sorted list of supported operator names."""
    return sorted(_OPERATORS.keys())


# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class PolicyV2Condition:
    """A single condition within a policy rule."""
    field: str
    operator: str
    value: Any

    def evaluate(self, context: dict[str, Any]) -> bool:
        """Evaluate this condition against a context dict."""
        actual = _resolve_field(context, self.field)
        if actual is None and self.operator not in ("eq", "neq"):
            return False
        return evaluate_operator(self.operator, actual, self.value)

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

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


@dataclass
class PolicyV2Rule:
    """A policy rule with conditions and an action."""
    rule_id: str
    name: str
    description: str
    conditions: list[PolicyV2Condition]
    action: PolicyV2Action
    severity: str = "MEDIUM"
    match_mode: str = "all"  # "all" (AND) or "any" (OR)
    enabled: bool = True
    tags: list[str] = field(default_factory=list)
    message: str = ""

    def evaluate(self, context: dict[str, Any]) -> bool:
        """Evaluate whether this rule's conditions match the context."""
        if not self.enabled:
            return False
        if not self.conditions:
            return False
        if self.match_mode == "any":
            return any(c.evaluate(context) for c in self.conditions)
        return all(c.evaluate(context) for c in self.conditions)

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

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PolicyV2Rule:
        return cls(
            rule_id=data.get("rule_id", ""),
            name=data.get("name", ""),
            description=data.get("description", ""),
            conditions=[
                PolicyV2Condition.from_dict(c)
                for c in data.get("conditions", [])
            ],
            action=PolicyV2Action(data.get("action", "BLOCK")),
            severity=data.get("severity", "MEDIUM"),
            match_mode=data.get("match_mode", "all"),
            enabled=data.get("enabled", True),
            tags=data.get("tags", []),
            message=data.get("message", ""),
        )


@dataclass
class PolicyV2RuleResult:
    """Result of evaluating a single rule against a finding."""
    rule_id: str
    rule_name: str
    action: PolicyV2Action
    verdict: PolicyV2Verdict
    severity: str
    message: str
    finding_id: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "rule_id": self.rule_id,
            "rule_name": self.rule_name,
            "action": self.action.value,
            "verdict": self.verdict.value,
            "severity": self.severity,
            "message": self.message,
            "finding_id": self.finding_id,
        }


@dataclass
class PolicyV2Result:
    """Aggregate result of evaluating a policy against findings."""
    blocked: list[PolicyV2RuleResult] = field(default_factory=list)
    warned: list[PolicyV2RuleResult] = field(default_factory=list)
    allowed: list[PolicyV2RuleResult] = field(default_factory=list)
    required_missing: list[PolicyV2RuleResult] = field(default_factory=list)
    passed: list[PolicyV2RuleResult] = field(default_factory=list)
    ci_pass: bool = True
    total_evaluated: int = 0
    timestamp: str = ""

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

    @property
    def block_count(self) -> int:
        return len(self.blocked)

    @property
    def warn_count(self) -> int:
        return len(self.warned)

    @property
    def allow_count(self) -> int:
        return len(self.allowed)

    @property
    def required_missing_count(self) -> int:
        return len(self.required_missing)

    def to_dict(self) -> dict[str, Any]:
        return {
            "ci_pass": self.ci_pass,
            "total_evaluated": self.total_evaluated,
            "block_count": self.block_count,
            "warn_count": self.warn_count,
            "allow_count": self.allow_count,
            "required_missing_count": self.required_missing_count,
            "blocked": [r.to_dict() for r in self.blocked],
            "warned": [r.to_dict() for r in self.warned],
            "allowed": [r.to_dict() for r in self.allowed],
            "required_missing": [r.to_dict() for r in self.required_missing],
            "passed": [r.to_dict() for r in self.passed],
            "timestamp": self.timestamp,
        }


@dataclass
class PolicyV2:
    """A complete policy with metadata and rules."""
    name: str = ""
    version: str = "1.0"
    description: str = ""
    level: str = "org"  # org, team, project
    parent: str = ""    # name of parent policy
    rules: list[PolicyV2Rule] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

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

    @property
    def enabled_rules(self) -> list[PolicyV2Rule]:
        return [r for r in self.rules if r.enabled]

    def get_rule(self, rule_id: str) -> PolicyV2Rule | None:
        """Get a rule by its ID."""
        for r in self.rules:
            if r.rule_id == rule_id:
                return r
        return None

    def add_rule(self, rule: PolicyV2Rule) -> None:
        """Add or replace a rule."""
        for i, existing in enumerate(self.rules):
            if existing.rule_id == rule.rule_id:
                self.rules[i] = rule
                return
        self.rules.append(rule)

    def remove_rule(self, rule_id: str) -> bool:
        """Remove a rule by ID. Returns True if removed."""
        for i, r in enumerate(self.rules):
            if r.rule_id == rule_id:
                self.rules.pop(i)
                return True
        return False

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "version": self.version,
            "description": self.description,
            "level": self.level,
            "parent": self.parent,
            "rules": [r.to_dict() for r in self.rules],
            "metadata": dict(self.metadata),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> PolicyV2:
        return cls(
            name=data.get("name", ""),
            version=data.get("version", "1.0"),
            description=data.get("description", ""),
            level=data.get("level", "org"),
            parent=data.get("parent", ""),
            rules=[PolicyV2Rule.from_dict(r) for r in data.get("rules", [])],
            metadata=data.get("metadata", {}),
        )

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

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


# ---------------------------------------------------------------------------
# Field resolution (supports dotted paths)
# ---------------------------------------------------------------------------

def _resolve_field(context: dict[str, Any], field_path: str) -> Any:
    """Resolve a dotted field path in a nested dict.

    E.g. "finding.severity" resolves context["finding"]["severity"].
    """
    parts = field_path.split(".")
    current: Any = context
    for part in parts:
        if isinstance(current, dict):
            current = current.get(part)
        else:
            return None
        if current is None:
            return None
    return current


# ---------------------------------------------------------------------------
# Policy evaluation
# ---------------------------------------------------------------------------

def _finding_to_context(finding: dict[str, Any]) -> dict[str, Any]:
    """Wrap a finding dict so fields are accessible for rule evaluation."""
    ctx = dict(finding)
    # Flatten nested fields for convenience
    if "severity" in finding:
        ctx["severity"] = finding["severity"]
    if "rule_id" in finding:
        ctx["rule_id"] = finding["rule_id"]
    return ctx


def evaluate_policy(
    findings: list[dict[str, Any]],
    policy: PolicyV2,
) -> PolicyV2Result:
    """Evaluate a list of findings against a policy.

    Each finding is a dict with keys like severity, rule_id, title, etc.
    Each enabled rule in the policy is evaluated against each finding.

    BLOCK verdicts cause ci_pass=False.
    REQUIRE rules that match NO finding produce REQUIRED_MISSING.

    Args:
        findings: List of finding dicts.
        policy: PolicyV2 to evaluate against.

    Returns:
        PolicyV2Result with categorized results.
    """
    result = PolicyV2Result()
    enabled = policy.enabled_rules

    # Track which REQUIRE rules matched at least once
    require_matched: dict[str, bool] = {}
    for rule in enabled:
        if rule.action == PolicyV2Action.REQUIRE:
            require_matched[rule.rule_id] = False

    for finding in findings:
        ctx = _finding_to_context(finding)
        finding_id = finding.get("rule_id", "") or finding.get("id", "")

        for rule in enabled:
            matched = rule.evaluate(ctx)
            result.total_evaluated += 1

            if rule.action == PolicyV2Action.ALLOW:
                if matched:
                    rr = PolicyV2RuleResult(
                        rule_id=rule.rule_id,
                        rule_name=rule.name,
                        action=rule.action,
                        verdict=PolicyV2Verdict.ALLOWED,
                        severity=rule.severity,
                        message=rule.message or f"Allowed by {rule.name}",
                        finding_id=finding_id,
                    )
                    result.allowed.append(rr)
                else:
                    result.passed.append(PolicyV2RuleResult(
                        rule_id=rule.rule_id,
                        rule_name=rule.name,
                        action=rule.action,
                        verdict=PolicyV2Verdict.PASSED,
                        severity=rule.severity,
                        message="",
                        finding_id=finding_id,
                    ))
            elif rule.action == PolicyV2Action.BLOCK:
                if matched:
                    rr = PolicyV2RuleResult(
                        rule_id=rule.rule_id,
                        rule_name=rule.name,
                        action=rule.action,
                        verdict=PolicyV2Verdict.BLOCKED,
                        severity=rule.severity,
                        message=rule.message or f"Blocked by {rule.name}",
                        finding_id=finding_id,
                    )
                    result.blocked.append(rr)
                    result.ci_pass = False
                else:
                    result.passed.append(PolicyV2RuleResult(
                        rule_id=rule.rule_id,
                        rule_name=rule.name,
                        action=rule.action,
                        verdict=PolicyV2Verdict.PASSED,
                        severity=rule.severity,
                        message="",
                        finding_id=finding_id,
                    ))
            elif rule.action == PolicyV2Action.WARN:
                if matched:
                    rr = PolicyV2RuleResult(
                        rule_id=rule.rule_id,
                        rule_name=rule.name,
                        action=rule.action,
                        verdict=PolicyV2Verdict.WARNED,
                        severity=rule.severity,
                        message=rule.message or f"Warning: {rule.name}",
                        finding_id=finding_id,
                    )
                    result.warned.append(rr)
                else:
                    result.passed.append(PolicyV2RuleResult(
                        rule_id=rule.rule_id,
                        rule_name=rule.name,
                        action=rule.action,
                        verdict=PolicyV2Verdict.PASSED,
                        severity=rule.severity,
                        message="",
                        finding_id=finding_id,
                    ))
            elif rule.action == PolicyV2Action.REQUIRE:
                if matched:
                    require_matched[rule.rule_id] = True
                    result.passed.append(PolicyV2RuleResult(
                        rule_id=rule.rule_id,
                        rule_name=rule.name,
                        action=rule.action,
                        verdict=PolicyV2Verdict.PASSED,
                        severity=rule.severity,
                        message="Requirement satisfied",
                        finding_id=finding_id,
                    ))

    # Check REQUIRE rules that never matched
    for rule in enabled:
        if rule.action == PolicyV2Action.REQUIRE:
            if not require_matched.get(rule.rule_id, False):
                rr = PolicyV2RuleResult(
                    rule_id=rule.rule_id,
                    rule_name=rule.name,
                    action=rule.action,
                    verdict=PolicyV2Verdict.REQUIRED_MISSING,
                    severity=rule.severity,
                    message=rule.message or f"Required: {rule.name} — not found",
                )
                result.required_missing.append(rr)
                result.ci_pass = False

    return result


# ---------------------------------------------------------------------------
# Policy merge (inheritance)
# ---------------------------------------------------------------------------

def merge_policies(parent: PolicyV2, child: PolicyV2) -> PolicyV2:
    """Merge a child policy into a parent, producing a combined policy.

    Inheritance semantics:
    - Child rules override parent rules with the same rule_id
    - Parent rules not overridden are inherited
    - Child metadata overrides parent metadata keys
    - The resulting policy has the child's name and level

    Args:
        parent: The parent (higher-level) policy.
        child: The child (lower-level) policy.

    Returns:
        A new PolicyV2 combining both.
    """
    # Index parent rules by rule_id
    parent_rules: dict[str, PolicyV2Rule] = {
        r.rule_id: r for r in parent.rules
    }

    # Child rules override
    child_rule_ids: set[str] = set()
    merged_rules: list[PolicyV2Rule] = []

    for rule in child.rules:
        child_rule_ids.add(rule.rule_id)
        merged_rules.append(rule)

    # Add inherited parent rules
    for rule_id, rule in parent_rules.items():
        if rule_id not in child_rule_ids:
            merged_rules.append(rule)

    # Merge metadata
    merged_meta = dict(parent.metadata)
    merged_meta.update(child.metadata)

    return PolicyV2(
        name=child.name or parent.name,
        version=child.version,
        description=child.description or parent.description,
        level=child.level,
        parent=parent.name,
        rules=merged_rules,
        metadata=merged_meta,
    )


def merge_policy_chain(policies: list[PolicyV2]) -> PolicyV2:
    """Merge a chain of policies from broadest (org) to narrowest (project).

    Args:
        policies: List of policies ordered from org to project.

    Returns:
        A single merged PolicyV2.
    """
    if not policies:
        return PolicyV2()
    result = policies[0]
    for policy in policies[1:]:
        result = merge_policies(result, policy)
    return result


# ---------------------------------------------------------------------------
# Policy validation
# ---------------------------------------------------------------------------

def validate_policy(policy: PolicyV2) -> list[str]:
    """Validate a policy for common errors.

    Returns a list of error messages. Empty list means valid.
    """
    errors: list[str] = []

    if not policy.name:
        errors.append("Policy must have a name")

    if policy.level not in ("org", "team", "project"):
        errors.append(f"Invalid level '{policy.level}', must be org/team/project")

    seen_ids: set[str] = set()
    for rule in policy.rules:
        if not rule.rule_id:
            errors.append("Rule missing rule_id")
        elif rule.rule_id in seen_ids:
            errors.append(f"Duplicate rule_id: {rule.rule_id}")
        else:
            seen_ids.add(rule.rule_id)

        if not rule.name:
            errors.append(f"Rule {rule.rule_id}: missing name")

        if not rule.conditions:
            errors.append(f"Rule {rule.rule_id}: no conditions defined")

        for cond in rule.conditions:
            if cond.operator not in _OPERATORS:
                errors.append(
                    f"Rule {rule.rule_id}: unsupported operator '{cond.operator}'"
                )
            if not cond.field:
                errors.append(f"Rule {rule.rule_id}: condition missing field")

        if rule.match_mode not in ("all", "any"):
            errors.append(
                f"Rule {rule.rule_id}: invalid match_mode '{rule.match_mode}'"
            )

    return errors


# ---------------------------------------------------------------------------
# Built-in policies
# ---------------------------------------------------------------------------

def get_builtin_org_policy() -> PolicyV2:
    """Return a built-in organizational security policy."""
    return PolicyV2(
        name="mcpaudit-org-default",
        version="1.0",
        description="Default organizational security policy for MCP servers",
        level="org",
        rules=[
            PolicyV2Rule(
                rule_id="ORG-001",
                name="Block critical findings",
                description="Block all CRITICAL severity findings",
                conditions=[
                    PolicyV2Condition("severity", "eq", "CRITICAL"),
                ],
                action=PolicyV2Action.BLOCK,
                severity="CRITICAL",
                message="Critical findings must be resolved before deployment",
            ),
            PolicyV2Rule(
                rule_id="ORG-002",
                name="Block high findings",
                description="Block HIGH severity findings",
                conditions=[
                    PolicyV2Condition("severity", "eq", "HIGH"),
                ],
                action=PolicyV2Action.BLOCK,
                severity="HIGH",
                message="High-severity findings must be resolved",
            ),
            PolicyV2Rule(
                rule_id="ORG-003",
                name="Warn on medium findings",
                description="Warn on MEDIUM severity findings",
                conditions=[
                    PolicyV2Condition("severity", "eq", "MEDIUM"),
                ],
                action=PolicyV2Action.WARN,
                severity="MEDIUM",
                message="Medium findings should be reviewed",
            ),
            PolicyV2Rule(
                rule_id="ORG-004",
                name="Block command injection",
                description="Block any command injection finding",
                conditions=[
                    PolicyV2Condition("rule_id", "matches", "^CMD-"),
                ],
                action=PolicyV2Action.BLOCK,
                severity="CRITICAL",
                message="Command injection must be remediated",
            ),
            PolicyV2Rule(
                rule_id="ORG-005",
                name="Block SQL injection",
                description="Block any SQL injection finding",
                conditions=[
                    PolicyV2Condition("rule_id", "matches", "^SQL-"),
                ],
                action=PolicyV2Action.BLOCK,
                severity="CRITICAL",
                message="SQL injection must be remediated",
            ),
            PolicyV2Rule(
                rule_id="ORG-006",
                name="Require authentication",
                description="At least one finding must indicate auth is present",
                conditions=[
                    PolicyV2Condition("title", "contains", "auth"),
                ],
                action=PolicyV2Action.REQUIRE,
                severity="HIGH",
                message="Authentication is required for all MCP servers",
            ),
        ],
    )


def get_builtin_team_policy() -> PolicyV2:
    """Return a sample team-level policy."""
    return PolicyV2(
        name="mcpaudit-team-default",
        version="1.0",
        description="Default team security policy",
        level="team",
        parent="mcpaudit-org-default",
        rules=[
            PolicyV2Rule(
                rule_id="TEAM-001",
                name="Warn on path traversal",
                description="Warn on path traversal findings",
                conditions=[
                    PolicyV2Condition("rule_id", "matches", "^PATH-"),
                ],
                action=PolicyV2Action.WARN,
                severity="HIGH",
                message="Path traversal detected — review and fix",
            ),
            PolicyV2Rule(
                rule_id="TEAM-002",
                name="Allow info findings",
                description="Whitelist INFO severity findings",
                conditions=[
                    PolicyV2Condition("severity", "eq", "INFO"),
                ],
                action=PolicyV2Action.ALLOW,
                severity="INFO",
                message="Info findings are accepted",
            ),
        ],
    )


# ---------------------------------------------------------------------------
# Reporting
# ---------------------------------------------------------------------------

def format_policy_result_text(result: PolicyV2Result) -> str:
    """Format a policy evaluation result as plain text."""
    lines: list[str] = []
    status = "PASS" if result.ci_pass else "FAIL"
    lines.append(f"Policy Evaluation: {status}")
    lines.append(f"Total rules evaluated: {result.total_evaluated}")
    lines.append(f"Blocked: {result.block_count}")
    lines.append(f"Warnings: {result.warn_count}")
    lines.append(f"Allowed: {result.allow_count}")
    lines.append(f"Required missing: {result.required_missing_count}")

    if result.blocked:
        lines.append("\n--- BLOCKED ---")
        for r in result.blocked:
            lines.append(f"  [{r.severity}] {r.rule_id}: {r.message}")

    if result.warned:
        lines.append("\n--- WARNINGS ---")
        for r in result.warned:
            lines.append(f"  [{r.severity}] {r.rule_id}: {r.message}")

    if result.required_missing:
        lines.append("\n--- REQUIRED MISSING ---")
        for r in result.required_missing:
            lines.append(f"  [{r.severity}] {r.rule_id}: {r.message}")

    return "\n".join(lines)


def format_policy_result_json(result: PolicyV2Result) -> str:
    """Format a policy evaluation result as JSON."""
    return json.dumps(result.to_dict(), indent=2)
