"""The inline action gate — ABOM's enforcement layer (the killer demo).

The strategy pivot in one module: stop *recording* what an agent did and start
*blocking* what it must not do. The gate sits at the tool-call boundary. Before
an action executes, it decides ALLOW or DENY against the agent's signed
Composition Manifest (the declared envelope), and **notarizes every decision** —
allowed or denied — into the Merkle transparency log under a signed tree head.

Design notes that matter to a buyer:
  * **Deny-by-default.** A tool not declared in the signed manifest is denied,
    not logged-and-allowed. That is the difference between a guardrail and a
    SIEM.
  * **Framework-agnostic.** ``Gate.check()`` is pure; ``@gated`` and
    ``guard_tool()`` adapt it to any Python agent without a framework dependency.
  * **Scope-honest.** The gate only claims to mediate what it actually sees (the
    tool calls routed through it). It does not pretend to catch egress it cannot
    observe — that boundary is declared in the manifest, not hand-waved.

A ``Decision`` is the thing handed to an auditor: the action, the verdict, the
rule that fired, and a Merkle inclusion proof anchoring it in the log.
"""
from __future__ import annotations

import fnmatch
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Callable

from .audit import canonical_json
from .merkle import MerkleLog, PersistentMerkleLog, verify_inclusion_hex
from .signer import LocalSigner, Signer, sign_payload


class Verdict(str, Enum):
    ALLOW = "allow"
    DENY = "deny"


class ActionDenied(Exception):
    """Raised when a gated action is denied. Carries the recorded Decision."""

    def __init__(self, decision: "Decision"):
        self.decision = decision
        super().__init__(
            f"ABOM gate DENIED action '{decision.action.get('tool')}': {decision.reason}"
        )


@dataclass
class Action:
    """A proposed agent action, normalized for the gate.

    ``kind`` selects what the action is checked against: ``"tool"`` (default)
    gates against the manifest's declared tools/MCP servers; ``"model"`` gates a
    model call against the manifest's declared models (used by runtime adapters).
    """
    tool: str
    args: dict = field(default_factory=dict)
    endpoint: str | None = None
    data_classification: str | None = None
    kind: str = "tool"

    def as_dict(self) -> dict:
        return {
            "tool": self.tool,
            "args": self.args,
            "endpoint": self.endpoint,
            "data_classification": self.data_classification,
            "kind": self.kind,
        }


@dataclass
class Decision:
    """The notarized outcome of a gate check — the artifact an auditor reads."""
    verdict: Verdict
    rule: str
    reason: str
    action: dict
    seq: int
    leaf_index: int
    created_at: str
    inclusion_proof: dict | None = None
    tree_head: dict | None = None
    head_signature: dict | None = None

    @property
    def allowed(self) -> bool:
        return self.verdict == Verdict.ALLOW

    def as_dict(self) -> dict:
        return {
            "verdict": self.verdict.value,
            "rule": self.rule,
            "reason": self.reason,
            "action": self.action,
            "seq": self.seq,
            "leaf_index": self.leaf_index,
            "created_at": self.created_at,
            "inclusion_proof": self.inclusion_proof,
            "tree_head": self.tree_head,
            "head_signature": self.head_signature,
        }


def _allowed_tools(manifest: dict) -> set[str]:
    return {
        c.get("name")
        for c in manifest.get("components", [])
        if c.get("type") in ("tool", "mcpServer") and c.get("name")
    }


def _declared_models(manifest: dict) -> set[str]:
    return {
        c.get("name")
        for c in manifest.get("components", [])
        if c.get("type") == "model" and c.get("name")
    }


def _tool_component(manifest: dict, tool: str) -> dict | None:
    for c in manifest.get("components", []):
        if c.get("type") in ("tool", "mcpServer") and c.get("name") == tool:
            return c
    return None


def evaluate(manifest: dict, action: Action) -> tuple[Verdict, str, str]:
    """Pure policy decision against the signed manifest. Deny-by-default.

    Returns (verdict, rule, reason). Order matters: the first failing rule wins,
    so the reason an auditor sees is the most specific violation.
    """
    # Model calls (from runtime adapters) gate against declared models, not tools.
    if getattr(action, "kind", "tool") == "model":
        models = _declared_models(manifest)
        if action.tool not in models:
            return (
                Verdict.DENY,
                "model_not_in_manifest",
                f"model '{action.tool}' is not in the signed Composition Manifest "
                f"(declared: {sorted(models) or 'none'})",
            )
        return (Verdict.ALLOW, "manifest_allows", "model is in the signed envelope")

    declared = _allowed_tools(manifest)

    # Rule 1: the tool must be declared in the signed manifest.
    if action.tool not in declared:
        return (
            Verdict.DENY,
            "tool_not_in_manifest",
            f"tool '{action.tool}' is not in the signed Composition Manifest "
            f"(declared: {sorted(declared) or 'none'})",
        )

    comp = _tool_component(manifest, action.tool) or {}

    # Rule 2: if the tool has an egress endpoint, it must be on its allowlist.
    if action.endpoint is not None:
        allowed_eps = comp.get("allowed_endpoints")
        if allowed_eps is not None and not any(
            fnmatch.fnmatch(action.endpoint, pat) for pat in allowed_eps
        ):
            return (
                Verdict.DENY,
                "endpoint_not_allowed",
                f"endpoint '{action.endpoint}' not in allowlist for tool "
                f"'{action.tool}' ({allowed_eps})",
            )

    # Rule 3: residency — confidential/restricted data must not egress.
    controls = manifest.get("controls", {})
    no_egress = controls.get("no_egress_classifications", ["confidential", "restricted"])
    if (
        action.data_classification in no_egress
        and comp.get("scope") == "egress"
    ):
        return (
            Verdict.DENY,
            "residency",
            f"'{action.data_classification}' data may not leave via egress tool "
            f"'{action.tool}'",
        )

    return (Verdict.ALLOW, "manifest_allows", "action is within the signed envelope")


class Gate:
    """Stateful gate over one agent run: evaluates + notarizes each action.

    Every decision is appended to a Merkle transparency log and the resulting
    tree head is signed, so the full set of decisions for a run is tamper-evident
    and each individual decision carries an inclusion proof.
    """

    def __init__(self, manifest: dict, *, signer: Signer | None = None, run_id: str = "run",
                 trusted_keys: set[str] | None = None, log_path=None):
        """Build a gate over ``manifest``.

        ``trusted_keys`` pins the signer(s) whose signature makes the manifest
        authentic. When given, the manifest signature is verified *before* the
        gate will enforce against it; a manifest that is unsigned, tampered, or
        re-signed by an untrusted key raises ``ValueError`` — the gate refuses to
        enforce an envelope it cannot authenticate. When omitted, the manifest is
        trusted as-is (unauthenticated, backward-compatible behaviour).

        ``log_path`` makes the transparency log durable: decisions are appended to
        an append-only file that survives process exit, so a later process (or an
        auditor) can reload the exact log and prove append-only growth. When
        omitted, the log is in-memory (backward-compatible).
        """
        if trusted_keys is not None:
            from . import sign as _sign
            if not _sign.verify_obj(manifest, trusted_keys):
                raise ValueError(
                    "refusing to gate against an untrusted manifest: signature is "
                    "missing, tampered, or signed by a key outside the trusted set"
                )
        self.manifest = manifest
        self.signer = signer or LocalSigner()
        self.run_id = run_id
        self.log = PersistentMerkleLog(log_path) if log_path is not None else MerkleLog()
        self.decisions: list[Decision] = []

    def check(self, action: Action | dict) -> Decision:
        """Evaluate an action, notarize the decision, and return it.

        Does NOT raise on deny — use ``enforce`` (or ``@gated``) for the
        block-the-call behaviour. ``check`` is for callers that want the verdict.
        """
        if isinstance(action, dict):
            action = Action(**action)
        verdict, rule, reason = evaluate(self.manifest, action)
        created_at = datetime.now(timezone.utc).isoformat()
        seq = len(self.decisions)

        entry = canonical_json(
            {
                "run_id": self.run_id,
                "seq": seq,
                "verdict": verdict.value,
                "rule": rule,
                "action": action.as_dict(),
                "created_at": created_at,
            }
        ).encode("utf-8")
        leaf_index = self.log.append(entry)

        head = self.log.head()
        head_sig = sign_payload(self.signer, canonical_json(head.as_dict()).encode("utf-8"))

        decision = Decision(
            verdict=verdict,
            rule=rule,
            reason=reason,
            action=action.as_dict(),
            seq=seq,
            leaf_index=leaf_index,
            created_at=created_at,
            inclusion_proof=self.log.prove_inclusion(leaf_index),
            tree_head=head.as_dict(),
            head_signature=head_sig,
        )
        self.decisions.append(decision)
        return decision

    def enforce(self, action: Action | dict) -> Decision:
        """Like ``check`` but raises ``ActionDenied`` if the action is denied."""
        decision = self.check(action)
        if not decision.allowed:
            raise ActionDenied(decision)
        return decision

    # -- adapters: wrap any Python callable / tool -------------------------
    def gated(self, tool_name: str | None = None,
              endpoint: Callable[..., str | None] | str | None = None):
        """Decorator: gate a tool function. Denied calls raise ``ActionDenied``.

        ``endpoint`` may be a static string or a callable(*args, **kwargs)->str
        to derive the egress target from the call.
        """
        def deco(fn: Callable) -> Callable:
            name = tool_name or fn.__name__

            def wrapper(*args, **kwargs):
                ep = endpoint(*args, **kwargs) if callable(endpoint) else endpoint
                self.enforce(Action(tool=name, args=kwargs, endpoint=ep))
                return fn(*args, **kwargs)

            wrapper.__name__ = getattr(fn, "__name__", name)
            wrapper.__wrapped__ = fn
            return wrapper

        return deco

    def guard_tool(self, fn: Callable, *, tool_name: str | None = None,
                   endpoint: str | None = None) -> Callable:
        """Wrap an existing callable (e.g. an MCP/LangChain tool) with the gate."""
        return self.gated(tool_name=tool_name, endpoint=endpoint)(fn)

    # -- verification (what an auditor runs) -------------------------------
    def head(self) -> dict:
        return self.log.head().as_dict()

    def consistency_since(self, old_size: int) -> dict:
        """RFC 6962 consistency proof from ``old_size`` to the current head — what
        a witness checks before cosigning that the log only grew (append-only)."""
        return self.log.prove_consistency(old_size)

    def verify(self) -> dict:
        """Re-verify every decision's inclusion proof against the current head.

        Returns {"valid": bool, "decisions": int, "broken_seq": int|None}.
        A third party can run the equivalent with only the public proofs + key.
        """
        for d in self.decisions:
            entry = canonical_json(
                {
                    "run_id": self.run_id,
                    "seq": d.seq,
                    "verdict": d.verdict.value,
                    "rule": d.rule,
                    "action": d.action,
                    "created_at": d.created_at,
                }
            ).encode("utf-8")
            proof = self.log.prove_inclusion(d.leaf_index)
            if not verify_inclusion_hex(entry, proof):
                return {"valid": False, "decisions": len(self.decisions), "broken_seq": d.seq}
        return {"valid": True, "decisions": len(self.decisions), "broken_seq": None}
