"""Small typed policy evaluator shared by management, router, and connectors."""

from __future__ import annotations

import fnmatch
import re
from collections.abc import Iterable, Mapping
from enum import StrEnum
from typing import Any
from uuid import UUID, uuid4

from pydantic import BaseModel, ConfigDict, Field, field_validator


class PolicyValidationError(ValueError):
    """A policy contains an unsupported or unsafe selector."""


class Decision(StrEnum):
    ALLOW = "allow"
    ALLOW_WITH_PREVIEW = "allow_with_preview"
    ALLOW_WITH_APPROVAL = "allow_with_approval"
    ALLOW_IF_CONDITIONS = "allow_if_conditions"
    ENFORCE = "enforce"
    DENY = "deny"
    DENY_AND_ALERT = "deny_and_alert"


_IDENTIFIER = re.compile(r"^[A-Za-z0-9._:/-]{1,256}$")
_MATCH_FIELDS = {
    "workspace",
    "workspaceId",
    "project",
    "projectId",
    "projects",
    "projectSlugs",
    "environment",
    "environments",
    "environmentClassifications",
    "client",
    "clientId",
    "clientIds",
    "principal",
    "principalType",
    "virtualServer",
    "virtualServers",
    "capability",
    "capabilities",
    "tool",
    "tools",
    "toolPatterns",
    "resourceType",
    "resourceTypes",
    "connector",
    "connectors",
    "riskClass",
    "riskClasses",
    "origin",
    "origins",
}
_CONDITION_FIELDS = {
    "requireEnabledClient",
    "requireBinding",
    "requireScope",
    "requireCapabilityPack",
    "requireConnectorBinding",
    "requireInputValidation",
    "requireOutputValidation",
    "requireRedaction",
    "maxArgumentBytes",
    "allEqual",
    "requiredClaims",
    "maxDefaultBytes",
    "maxDefaultItems",
    "redactionProfile",
}


class InvocationContext(BaseModel):
    """Validated identity and target context used at every sensitive boundary."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    workspace: str = Field(min_length=1, max_length=128)
    project: str = Field(min_length=1, max_length=128)
    environment: str = Field(pattern=r"^[a-z0-9][a-z0-9-]{0,31}$")
    client_id: UUID | None = None
    principal: str = Field(min_length=1, max_length=256)
    principal_type: str = Field(default="machine", max_length=64)
    virtual_server: str | None = Field(default=None, max_length=128)
    capability: str | None = Field(default=None, max_length=256)
    tool: str = Field(min_length=1, max_length=256)
    connector: str = Field(min_length=1, max_length=64)
    risk_class: str = Field(default="R1", pattern=r"^(?:R[0-9]+|RX)$")
    origin: str = Field(default="mcp", max_length=64)
    resource_type: str | None = Field(default=None, max_length=64)
    environment_classification: str | None = Field(default=None, max_length=64)
    binding_present: bool = False
    client_enabled: bool = False
    arguments: Mapping[str, Any] = Field(default_factory=dict)
    trace_id: str = Field(default_factory=lambda: uuid4().hex, pattern=r"^[A-Za-z0-9._:-]{1,128}$")


class PolicyRule(BaseModel):
    """A catalog or persisted rule with only typed selector primitives."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    slug: str = Field(pattern=r"^[a-z0-9][a-z0-9-]{0,62}$")
    workspace: str | None = None
    project: str | None = None
    priority: int = Field(ge=0, le=100000)
    description: str = Field(min_length=1, max_length=1000)
    match: dict[str, Any] = Field(default_factory=dict)
    conditions: dict[str, Any] = Field(default_factory=dict)
    decision: Decision
    terminal: bool = False
    source: str = Field(default="catalog", max_length=64)
    enabled: bool = True

    @field_validator("match", "conditions")
    @classmethod
    def safe_maps(cls, value: dict[str, Any]) -> dict[str, Any]:
        if len(value) > 32:
            raise PolicyValidationError("policy selector map is too large")
        return value

    @classmethod
    def from_document(cls, document: Mapping[str, Any]) -> PolicyRule:
        metadata = document.get("metadata") or {}
        spec = document.get("spec") or {}
        if not isinstance(metadata, Mapping) or not isinstance(spec, Mapping):
            raise PolicyValidationError("policy document metadata and spec must be objects")
        match = dict(spec.get("match") or {})
        conditions = dict(spec.get("conditions") or {})
        unknown_match = set(match) - _MATCH_FIELDS
        unknown_conditions = set(conditions) - _CONDITION_FIELDS
        if unknown_match or unknown_conditions:
            raise PolicyValidationError(
                f"unsupported policy fields: {sorted(unknown_match | unknown_conditions)}"
            )
        project = spec.get("project")
        workspace = spec.get("workspace")
        return cls(
            slug=str(metadata.get("slug", "")),
            workspace=str(workspace) if workspace is not None else None,
            project=str(project) if project is not None else None,
            priority=int(spec.get("priority", 0)),
            description=str(spec.get("description", "")),
            match=match,
            conditions=conditions,
            decision=Decision(str(spec.get("decision", "deny"))),
            terminal=bool(spec.get("terminal", False)),
            source=str(spec.get("source", "catalog")),
            enabled=bool(spec.get("enabled", True)),
        )


class TraceStep(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    phase: str
    rule: str | None = None
    matched: bool
    reason: str


class DecisionTrace(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    trace_id: str
    decision: Decision
    reason_code: str
    matched_rule_ids: tuple[str, ...] = ()
    steps: tuple[TraceStep, ...] = ()


def _values(value: Any) -> tuple[str, ...]:
    if isinstance(value, list | tuple | set):
        return tuple(str(item) for item in value)
    return (str(value),)


def _context_value(context: InvocationContext, key: str) -> tuple[str, ...]:
    mapping: dict[str, Any] = {
        "workspace": context.workspace,
        "workspaceId": context.workspace,
        "project": context.project,
        "projectId": context.project,
        "projectSlugs": context.project,
        "projects": context.project,
        "environment": context.environment,
        "environments": context.environment,
        "environmentClassifications": context.environment_classification or context.environment,
        "resourceType": context.resource_type or "",
        "resourceTypes": context.resource_type or "",
        "client": str(context.client_id) if context.client_id else "",
        "clientId": str(context.client_id) if context.client_id else "",
        "clientIds": str(context.client_id) if context.client_id else "",
        "principal": context.principal,
        "principalType": context.principal_type,
        "virtualServer": context.virtual_server or "",
        "capability": context.capability or "",
        "capabilities": context.capability or "",
        "tool": context.tool,
        "tools": context.tool,
        "toolPatterns": context.tool,
        "connector": context.connector,
        "connectors": context.connector,
        "riskClass": context.risk_class,
        "riskClasses": context.risk_class,
        "origin": context.origin,
        "origins": context.origin,
    }
    return _values(mapping.get(key, ""))


def _rule_scope_matches(rule: PolicyRule, context: InvocationContext) -> bool:
    return (rule.workspace is None or rule.workspace == context.workspace) and (
        rule.project is None or rule.project == context.project
    )


def _matches(rule: PolicyRule, context: InvocationContext) -> bool:
    if not _rule_scope_matches(rule, context):
        return False
    for key, expected in rule.match.items():
        if key not in _MATCH_FIELDS:
            raise PolicyValidationError(f"unsupported policy match field: {key}")
        actual = _context_value(context, key)
        expected_values = _values(expected)
        if key in {"toolPatterns"}:
            if not any(
                fnmatch.fnmatchcase(value, pattern)
                for value in actual
                for pattern in expected_values
            ):
                return False
        elif not any(value in expected_values for value in actual):
            return False
    return True


def _conditions_match(rule: PolicyRule, context: InvocationContext) -> bool:
    for key, expected in rule.conditions.items():
        if key not in _CONDITION_FIELDS:
            raise PolicyValidationError(f"unsupported policy condition: {key}")
        if key == "requireEnabledClient" and bool(expected) and not context.client_enabled:
            return False
        if key == "requireBinding" and bool(expected) and not context.binding_present:
            return False
        if key == "maxArgumentBytes" and len(str(dict(context.arguments)).encode()) > int(expected):
            return False
        if key in {
            "allEqual",
            "requiredClaims",
            "maxDefaultBytes",
            "maxDefaultItems",
            "redactionProfile",
        }:
            return False
    return True


class PolicyEngine:
    """Deterministic precedence: validation, global deny, explicit deny, allow, default deny."""

    def __init__(self, policies: Iterable[PolicyRule]) -> None:
        self.policies = tuple(
            sorted(
                (item for item in policies if item.enabled),
                key=lambda item: (-item.priority, item.slug),
            )
        )
        for policy in self.policies:
            unknown = set(policy.match) - _MATCH_FIELDS
            if unknown:
                raise PolicyValidationError(f"unsupported policy match fields: {sorted(unknown)}")
            if (
                policy.decision
                in {
                    Decision.ALLOW,
                    Decision.ALLOW_IF_CONDITIONS,
                    Decision.ALLOW_WITH_APPROVAL,
                    Decision.ALLOW_WITH_PREVIEW,
                }
                and policy.project is None
                and policy.slug != "global-deny"
            ):
                # Workspace-wide allows must not expand a project scope implicitly.
                continue

    def evaluate(self, context: InvocationContext) -> DecisionTrace:
        steps: list[TraceStep] = []
        if not context.binding_present:
            return DecisionTrace(
                trace_id=context.trace_id,
                decision=Decision.DENY,
                reason_code="MISSING_BINDING",
                steps=(
                    TraceStep(
                        phase="identity", matched=False, reason="client/project binding is required"
                    ),
                ),
            )
        for rule in self.policies:
            matched = _matches(rule, context)
            steps.append(
                TraceStep(
                    phase="global-deny" if rule.project is None else "match",
                    rule=rule.slug,
                    matched=matched,
                    reason="selector matched" if matched else "selector did not match",
                )
            )
            if (
                matched
                and rule.project is None
                and rule.decision in {Decision.DENY, Decision.DENY_AND_ALERT}
                and rule.terminal
            ):
                return self._result(context, rule.decision, "GLOBAL_DENY", [rule.slug], steps)
        for rule in self.policies:
            if (
                rule.slug == "global-deny"
                or rule.decision not in {Decision.DENY, Decision.DENY_AND_ALERT}
                or not _matches(rule, context)
            ):
                continue
            return self._result(context, rule.decision, "EXPLICIT_DENY", [rule.slug], steps)
        for rule in self.policies:
            if rule.decision not in {
                Decision.ALLOW,
                Decision.ALLOW_IF_CONDITIONS,
                Decision.ALLOW_WITH_APPROVAL,
                Decision.ALLOW_WITH_PREVIEW,
            } or not _matches(rule, context):
                continue
            conditions = _conditions_match(rule, context)
            steps.append(
                TraceStep(
                    phase="conditions",
                    rule=rule.slug,
                    matched=conditions,
                    reason="conditions satisfied" if conditions else "conditions not satisfied",
                )
            )
            if conditions:
                return self._result(context, rule.decision, "ALLOW_RULE", [rule.slug], steps)
        return self._result(context, Decision.DENY, "DEFAULT_DENY", [], steps)

    @staticmethod
    def _result(
        context: InvocationContext,
        decision: Decision,
        reason: str,
        rules: list[str],
        steps: list[TraceStep],
    ) -> DecisionTrace:
        return DecisionTrace(
            trace_id=context.trace_id,
            decision=decision,
            reason_code=reason,
            matched_rule_ids=tuple(rules),
            steps=tuple(steps),
        )


class PolicySimulator:
    """Runs the exact engine against a known-entity inventory without side effects."""

    def __init__(
        self,
        engine: PolicyEngine,
        *,
        known_projects: Iterable[str] = (),
        known_clients: Iterable[UUID] = (),
        known_tools: Iterable[str] = (),
    ) -> None:
        self.engine = engine
        self.known_projects = frozenset(known_projects)
        self.known_clients = frozenset(known_clients)
        self.known_tools = frozenset(known_tools)

    def simulate(self, context: InvocationContext) -> DecisionTrace:
        if self.known_projects and context.project not in self.known_projects:
            return DecisionTrace(
                trace_id=context.trace_id,
                decision=Decision.DENY,
                reason_code="UNKNOWN_PROJECT",
                steps=(
                    TraceStep(
                        phase="simulation", matched=False, reason="project is not a known entity"
                    ),
                ),
            )
        if self.known_tools and context.tool not in self.known_tools:
            return DecisionTrace(
                trace_id=context.trace_id,
                decision=Decision.DENY,
                reason_code="UNKNOWN_TOOL",
                steps=(
                    TraceStep(
                        phase="simulation", matched=False, reason="tool is not a known entity"
                    ),
                ),
            )
        if context.client_id and self.known_clients and context.client_id not in self.known_clients:
            return DecisionTrace(
                trace_id=context.trace_id,
                decision=Decision.DENY,
                reason_code="UNKNOWN_CLIENT",
                steps=(
                    TraceStep(
                        phase="simulation", matched=False, reason="client is not a known entity"
                    ),
                ),
            )
        return self.engine.evaluate(context)