"""ActionGate LangChain Integration — gate tool execution with ActionGate primitives.

Two integration patterns:

1. @gated_tool decorator (simplest):

    from actiongate_langchain import gated_tool
    from actiongate import Engine, Gate, Policy

    ag = Engine()

    @gated_tool(
        actiongate=ag,
        gate=Gate("tools", "search", "agent:1"),
        policy=Policy(max_calls=10, window=60),
    )
    def search_web(query: str) -> str:
        '''Search the web for a query.'''
        return api.search(query)

    # Use directly with LangChain agents
    agent = create_react_agent(llm, [search_web], prompt)

2. GatedTool class (full control):

    from actiongate_langchain import GatedTool

    tool = GatedTool.from_function(
        func=search_fn,
        name="search_web",
        description="Search the web",
        actiongate=ag,
        gate=Gate("tools", "search", "agent:1"),
        policy=Policy(max_calls=10, window=60),
    )
"""

from __future__ import annotations

import json
from decimal import Decimal
from functools import wraps
from typing import Any, Callable, Optional, Type, TYPE_CHECKING

from langchain_core.tools import BaseTool, ToolException
from pydantic import BaseModel, ConfigDict, PrivateAttr

if TYPE_CHECKING:
    from actiongate import Engine as ActionEngine, Gate, Policy
    from budgetgate import Engine as BudgetEngine, Ledger, Budget
    from rulegate import Engine as RuleEngine, Rule, Ruleset
    from auditgate import Engine as AuditEngine, Trail, AuditPolicy


class GatedTool(BaseTool):
    """A LangChain tool that runs through ActionGate primitives before execution.

    Gates evaluate in order: ActionGate → BudgetGate → RuleGate → execute → AuditGate.
    First block raises ToolException with the gate's block message.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    # Tool metadata
    name: str = ""
    description: str = ""
    args_schema: Optional[Type[BaseModel]] = None

    # Private attributes — not Pydantic model fields
    _fn: Callable[..., Any] = PrivateAttr(default=None)
    _ag: Optional[Any] = PrivateAttr(default=None)
    _bg: Optional[Any] = PrivateAttr(default=None)
    _rg: Optional[Any] = PrivateAttr(default=None)
    _audit: Optional[Any] = PrivateAttr(default=None)
    _gate: Optional[Any] = PrivateAttr(default=None)
    _policy: Optional[Any] = PrivateAttr(default=None)
    _ledger: Optional[Any] = PrivateAttr(default=None)
    _budget: Optional[Any] = PrivateAttr(default=None)
    _cost: Optional[Decimal] = PrivateAttr(default=None)
    _rule: Optional[Any] = PrivateAttr(default=None)
    _ruleset: Optional[Any] = PrivateAttr(default=None)
    _trail: Optional[Any] = PrivateAttr(default=None)
    _audit_policy: Optional[Any] = PrivateAttr(default=None)

    def _run(self, **kwargs: Any) -> Any:
        """Run the tool through the gate pipeline."""
        decisions: dict[str, Any] = {}

        # ActionGate: rate limiting
        if self._ag and self._gate:
            decision = self._ag.check(self._gate, self._policy)
            decisions["actiongate"] = decision.to_dict()
            if decision.blocked:
                self._audit_log("BLOCK", "actiongate", decisions)
                raise ToolException(f"[actiongate] {decision.message}")

        # BudgetGate: spend control
        if self._bg and self._ledger and self._cost is not None:
            decision = self._bg.check(self._ledger, self._cost, self._budget)
            decisions["budgetgate"] = decision.to_dict()
            if decision.blocked:
                self._audit_log("BLOCK", "budgetgate", decisions)
                raise ToolException(f"[budgetgate] {decision.message}")

        # RuleGate: policy enforcement
        if self._rg and self._rule and self._ruleset:
            decision = self._rg.check(self._rule, self._ruleset, kwargs=kwargs)
            decisions["rulegate"] = decision.to_dict()
            if decision.blocked:
                self._audit_log("BLOCK", "rulegate", decisions)
                raise ToolException(f"[rulegate] {decision.message}")

        # Execute
        result = self._fn(**kwargs)

        # AuditGate: log success
        self._audit_log("ALLOW", "pipeline", decisions)

        return result

    def _audit_log(self, verdict_str: str, gate_type: str,
                   decisions: dict[str, Any]) -> None:
        if self._audit is None or self._trail is None:
            return
        try:
            from auditgate import Verdict, Severity
            verdict_map = {"ALLOW": Verdict.ALLOW, "BLOCK": Verdict.BLOCK}
            severity_map = {"ALLOW": Severity.INFO, "BLOCK": Severity.WARN}
            self._audit.record(
                self._trail,
                verdict=verdict_map.get(verdict_str, Verdict.ERROR),
                severity=severity_map.get(verdict_str, Severity.ERROR),
                gate_type=gate_type,
                gate_identity=str(self._trail),
                reason=f"{verdict_str} by {gate_type}",
                detail=decisions,
                policy=self._audit_policy,
            )
        except Exception:
            pass

    @classmethod
    def from_function(
        cls,
        func: Callable[..., Any],
        name: str,
        description: str,
        *,
        args_schema: Optional[Type[BaseModel]] = None,
        actiongate: "ActionEngine | None" = None,
        gate: "Gate | None" = None,
        policy: "Policy | None" = None,
        budgetgate: "BudgetEngine | None" = None,
        ledger: "Ledger | None" = None,
        budget: "Budget | None" = None,
        cost: Decimal | None = None,
        rulegate: "RuleEngine | None" = None,
        rule: "Rule | None" = None,
        ruleset: "Ruleset | None" = None,
        auditgate: "AuditEngine | None" = None,
        trail: "Trail | None" = None,
        audit_policy: "AuditPolicy | None" = None,
    ) -> "GatedTool":
        """Create a GatedTool from a function with gate configurations."""
        tool = cls(
            name=name,
            description=description,
            args_schema=args_schema,
        )
        tool._fn = func
        tool._ag = actiongate
        tool._gate = gate
        tool._policy = policy
        tool._bg = budgetgate
        tool._ledger = ledger
        tool._budget = budget
        tool._cost = cost
        tool._rg = rulegate
        tool._rule = rule
        tool._ruleset = ruleset
        tool._audit = auditgate
        tool._trail = trail
        tool._audit_policy = audit_policy
        return tool


def gated_tool(
    *,
    actiongate: "ActionEngine | None" = None,
    gate: "Gate | None" = None,
    policy: "Policy | None" = None,
    budgetgate: "BudgetEngine | None" = None,
    ledger: "Ledger | None" = None,
    budget: "Budget | None" = None,
    cost: Decimal | None = None,
    rulegate: "RuleEngine | None" = None,
    rule: "Rule | None" = None,
    ruleset: "Ruleset | None" = None,
    auditgate: "AuditEngine | None" = None,
    trail: "Trail | None" = None,
    audit_policy: "AuditPolicy | None" = None,
    args_schema: Optional[Type[BaseModel]] = None,
) -> Callable[[Callable[..., Any]], GatedTool]:
    """Decorator that converts a function into a gated LangChain tool.

    Uses the function's name and docstring for the tool's name and description.

    Example:
        @gated_tool(
            actiongate=ag,
            gate=Gate("tools", "search", "agent:1"),
            policy=Policy(max_calls=10, window=60),
        )
        def search_web(query: str) -> str:
            '''Search the web for a query.'''
            return api.search(query)
    """
    def decorator(fn: Callable[..., Any]) -> GatedTool:
        return GatedTool.from_function(
            func=fn,
            name=fn.__name__,
            description=fn.__doc__ or fn.__name__,
            args_schema=args_schema,
            actiongate=actiongate, gate=gate, policy=policy,
            budgetgate=budgetgate, ledger=ledger, budget=budget, cost=cost,
            rulegate=rulegate, rule=rule, ruleset=ruleset,
            auditgate=auditgate, trail=trail, audit_policy=audit_policy,
        )
    return decorator
