"""
eaigins_sdk.decorators
──────────────────────
Zero-boilerplate governance via Python decorators.

The ``@governed_tool`` decorator wraps any callable so that every call
is intercepted by Eaigins before the function body executes.

Usage — standalone function::

    from eaigins_sdk.decorators import governed_tool
    from eaigins_sdk import ControlPlane
    from eaigins_sdk.policies import AML_POLICIES

    cp = ControlPlane(policies=AML_POLICIES)

    @governed_tool(
        cp,
        tool_name       = "export_customer_data",
        context_sources = ["crm_customer_profile", "payment_records"],
        agent_id        = "my-agent-v1",
    )
    def export_customer_data(scope: str, include_pii: bool) -> str:
        # Only reached if Eaigins ALLOWs
        return call_real_api(scope, include_pii)

    result = export_customer_data("ALL", True)   # raises GovernanceViolationError

Usage — dynamic params (tool_params resolved at call time)::

    @governed_tool(
        cp,
        tool_name = "process_refund",
        context_sources = ["payment_records"],
        resolve_params = lambda args, kwargs: {
            "amount":      kwargs.get("amount", 0),
            "customer_id": kwargs.get("customer_id", ""),
        },
    )
    def process_refund(customer_id: str, amount: float) -> dict:
        ...

Usage — LangChain @tool compatible::

    from langchain_core.tools import tool as lc_tool

    @lc_tool
    @governed_tool(cp, tool_name="get_patient_record", context_sources=["patient_ehr"])
    def get_patient_record(patient_id: str, session_id: str) -> str:
        ...
"""

from __future__ import annotations

import functools
from typing import Any, Callable, Optional

from .models         import AgentAction, GovernanceViolationError
from .control_plane  import ControlPlane


def governed_tool(
    control_plane:   ControlPlane,
    *,
    tool_name:       Optional[str]                                     = None,
    context_sources: Optional[list[str]]                               = None,
    agent_id:        str                                               = "sdk-agent",
    reasoning:       str                                               = "",
    metadata:        Optional[dict]                                    = None,
    resolve_params:  Optional[Callable[[tuple, dict], dict]]           = None,
    raise_on_escalate: bool = False,
) -> Callable:
    """
    Decorator that wraps a function with Eaigins governance.

    Args:
        control_plane:      The ``ControlPlane`` instance to evaluate against.
        tool_name:          The enterprise tool name for policy matching.
                            Defaults to the decorated function's name.
        context_sources:    RAG/retrieval sources this tool accesses.
        agent_id:           Agent identifier stamped on the ``AgentAction``.
        reasoning:          Static reasoning string. Override per-call via
                            ``_eaigins_reasoning`` keyword argument.
        metadata:           Static metadata dict (e.g. regulatory_framework).
        resolve_params:     ``(args, kwargs) -> dict`` callable that extracts
                            tool_params from the function's call arguments.
                            If None, ``kwargs`` is used directly.
        raise_on_escalate:  If True, also raise on ESCALATE (not just BLOCK).
    """
    def decorator(func: Callable) -> Callable:
        effective_tool_name = tool_name or func.__name__

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Allow per-call overrides via private kwargs
            call_reasoning = kwargs.pop("_eaigins_reasoning", reasoning)
            call_agent_id  = kwargs.pop("_eaigins_agent_id",  agent_id)
            call_prompt    = kwargs.pop("_eaigins_prompt",
                                        f"Invoking {effective_tool_name}")

            # Resolve tool_params
            if resolve_params is not None:
                params = resolve_params(args, kwargs)
            else:
                params = {k: v for k, v in kwargs.items()
                          if not k.startswith("_eaigins_")}

            action = AgentAction(
                agent_id        = call_agent_id,
                prompt          = call_prompt,
                tool_name       = effective_tool_name,
                tool_params     = params,
                context_sources = context_sources or [],
                reasoning       = call_reasoning,
                metadata        = metadata or {},
            )

            decision = control_plane.evaluate(action)

            from .models import ExecutionDecision
            should_raise = (
                decision.decision == ExecutionDecision.BLOCK or
                (raise_on_escalate and decision.decision == ExecutionDecision.ESCALATE)
            )
            if should_raise:
                raise GovernanceViolationError(decision)

            return func(*args, **kwargs)

        # Attach the control plane reference for introspection
        wrapper._eaigins_cp        = control_plane
        wrapper._eaigins_tool_name = effective_tool_name
        return wrapper

    return decorator


class GovernedMixin:
    """
    Mixin for agent classes that want to evaluate actions declaratively.

    Usage::

        class MyAgent(GovernedMixin):
            def __init__(self):
                self.control_plane = ControlPlane(policies=AML_POLICIES)
                self.agent_id      = "my-agent-v1"

            def do_something(self, customer_id: str):
                self.intercept(
                    tool_name       = "get_customer_full_profile",
                    tool_params     = {"customer_id": customer_id},
                    context_sources = ["crm_customer_profile"],
                    prompt          = f"Retrieve profile for {customer_id}",
                )
                # Only reached if ALLOW
                return self._real_get_profile(customer_id)
    """

    control_plane: ControlPlane
    agent_id: str = "sdk-agent"

    def intercept(
        self,
        tool_name:       str,
        tool_params:     Optional[dict]       = None,
        context_sources: Optional[list[str]]  = None,
        prompt:          str                  = "",
        reasoning:       str                  = "",
        metadata:        Optional[dict]       = None,
        raise_on_escalate: bool               = False,
    ) -> None:
        """
        Evaluate an action and raise ``GovernanceViolationError`` if blocked.
        Returns ``None`` if ALLOW.
        """
        action = AgentAction(
            agent_id        = getattr(self, "agent_id", "sdk-agent"),
            prompt          = prompt or f"Invoking {tool_name}",
            tool_name       = tool_name,
            tool_params     = tool_params or {},
            context_sources = context_sources or [],
            reasoning       = reasoning,
            metadata        = metadata or {},
        )
        decision = self.control_plane.evaluate(action)

        from .models import ExecutionDecision
        if decision.decision == ExecutionDecision.BLOCK:
            raise GovernanceViolationError(decision)
        if raise_on_escalate and decision.decision == ExecutionDecision.ESCALATE:
            raise GovernanceViolationError(decision)
