"""
AgentX integration hooks for Agendex governance.

This module provides framework-agnostic hooks that can be wired into
manager/assistant/tool execution boundaries. The goal is to keep client
integration small while still capturing rich trace context.
"""
from __future__ import annotations

import contextvars
import importlib
import inspect
import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from functools import wraps
from time import perf_counter
from typing import Any, Callable, Dict, Mapping, MutableMapping, Optional, Type
from uuid import uuid4

from ..client import AgendexClient
from ..errors import AgendexError, DeniedError, PendingApprovalError

logger = logging.getLogger("agendex.integrations.agentx")

_ACTIVE_AGENTX_STATE: contextvars.ContextVar[Any] = contextvars.ContextVar(
    "agendex_agentx_active_state", default=None
)
_DEFAULT_AGENTX_HOOKS: Optional["AgentXGovernanceHooks"] = None


def _utc_now_iso() -> str:
    return datetime.now(tz=timezone.utc).isoformat()


def _safe_text(value: Any, *, max_chars: int = 1200) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        text = value
    else:
        try:
            text = json.dumps(value, ensure_ascii=False)
        except Exception:
            text = str(value)
    if len(text) <= max_chars:
        return text
    return text[:max_chars].rstrip() + "..."


def _safe_json(value: Any, *, max_chars: int = 1200) -> Any:
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, Mapping):
        return {str(k): _safe_json(v, max_chars=max_chars) for k, v in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [_safe_json(v, max_chars=max_chars) for v in value]
    if hasattr(value, "model_dump"):
        try:
            return _safe_json(value.model_dump(), max_chars=max_chars)
        except Exception:
            pass
    if hasattr(value, "dict"):
        try:
            return _safe_json(value.dict(), max_chars=max_chars)
        except Exception:
            pass
    return _safe_text(value, max_chars=max_chars)


def _method_references_name(func: Any, name: str) -> bool:
    code = getattr(func, "__code__", None)
    if not code:
        return False
    names = getattr(code, "co_names", ())
    return name in names


def _extract_assistant_response_from_state(state: Any) -> Any:
    if not isinstance(state, MutableMapping):
        return None
    messages = state.get("assistant_messages_current_turn")
    if not isinstance(messages, list) or not messages:
        return None
    for msg in reversed(messages):
        if hasattr(msg, "tool_calls"):
            return msg
        if getattr(msg, "type", None) == "ai":
            return msg
    return messages[-1]


def _extract_tool_outputs_from_state(state: Any) -> list[Any]:
    if not isinstance(state, MutableMapping):
        return []
    messages = state.get("assistant_messages_current_turn")
    if not isinstance(messages, list) or not messages:
        return []
    outputs: list[Any] = []
    for msg in messages:
        if getattr(msg, "type", None) == "tool":
            outputs.append(getattr(msg, "content", msg))
    return outputs


@dataclass
class AgentXGovernanceConfig:
    """Configuration for AgentX governance hooks."""

    # Governance defaults
    task: str = "agentx_workflow"
    task_state_key: str = "task_name"
    action_prefix: str = "tool."
    action_map: Dict[str, str] = field(default_factory=dict)

    # State keys
    trace_state_key: str = "agendex_trace_id"
    trace_context_key: str = "trace_id"
    user_prompt_state_key: str = "user_query"
    manager_reasoning_state_key: str = "manager_reasoning"
    manager_decision_state_key: str = "manager_routing_decision"
    assistant_state_key: str = "last_assistant_routed"
    assistant_messages_state_key: str = "assistant_messages_current_turn"
    user_id_state_keys: tuple[str, ...] = ("user_id", "user_email", "operator_id")
    session_id_state_keys: tuple[str, ...] = ("session_id", "chat_id", "thread_id", "conversation_id")

    # Behavior
    persist_events: bool = True
    capture_message_preview: bool = False
    fallback_on_transport_error: bool = False
    reasoning_preview_chars: int = 1200
    message_preview_chars: int = 512

    # Optional callback for local UI/debug
    on_event: Optional[Callable[[Dict[str, Any]], None]] = None


class AgentXGovernanceHooks:
    """
    Governance hooks for AgentX manager/assistant/tool boundaries.

    The hooks are intentionally lightweight and optional. If they are not wired,
    AgentX behavior is unchanged.
    """

    _LAST_SPAN_KEY = "_agendex_last_span_id"
    _MANAGER_SPAN_KEY = "_agendex_manager_span_id"
    _ASSISTANT_SPAN_PREFIX = "_agendex_assistant_span::"

    def __init__(
        self,
        agendex: AgendexClient,
        config: Optional[AgentXGovernanceConfig] = None,
    ) -> None:
        self.agendex = agendex
        self.config = config or AgentXGovernanceConfig()

    def _assistant_span_key(self, assistant_name: str) -> str:
        return f"{self._ASSISTANT_SPAN_PREFIX}{assistant_name}"

    @staticmethod
    def _state_get(state: Any, key: str, default: Any = None) -> Any:
        if isinstance(state, MutableMapping):
            return state.get(key, default)
        return default

    @staticmethod
    def _state_set(state: Any, key: str, value: Any) -> None:
        if isinstance(state, MutableMapping):
            state[key] = value

    def _resolve_task(self, state: Any) -> str:
        task = self._state_get(state, self.config.task_state_key)
        if isinstance(task, str) and task.strip():
            return task.strip()
        return self.config.task

    def _ensure_trace_id(self, state: Any) -> str:
        trace_id = self._state_get(state, self.config.trace_context_key) or self._state_get(
            state, self.config.trace_state_key
        )
        if not trace_id:
            trace_id = uuid4().hex
            self._state_set(state, self.config.trace_state_key, trace_id)
        self._state_set(state, self.config.trace_context_key, trace_id)
        return str(trace_id)

    def _new_span_id(self) -> str:
        return uuid4().hex

    def _extract_assistant_reasoning(self, state: Any) -> Optional[str]:
        messages = self._state_get(state, self.config.assistant_messages_state_key)
        if not isinstance(messages, list) or not messages:
            return None
        last = messages[-1]
        content = getattr(last, "content", None)
        if content is None:
            return None
        text = _safe_text(content, max_chars=self.config.reasoning_preview_chars)
        return text or None

    def _message_preview(self, messages: Any) -> Optional[list[str]]:
        if not self.config.capture_message_preview:
            return None
        if not isinstance(messages, list):
            return None
        preview: list[str] = []
        for msg in messages[-5:]:
            content = getattr(msg, "content", msg)
            text = _safe_text(content, max_chars=self.config.message_preview_chars)
            if text:
                preview.append(text)
        return preview or None

    def _first_state_string(self, state: Any, keys: tuple[str, ...]) -> Optional[str]:
        for key in keys:
            value = self._state_get(state, key)
            if value is None:
                continue
            text = str(value).strip()
            if text:
                return text
        return None

    def _build_context(
        self,
        *,
        state: Any,
        reasoning: Optional[str] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        context: Dict[str, Any] = {}
        user_id = self._first_state_string(state, self.config.user_id_state_keys)
        if user_id:
            context["user_id"] = user_id
        session_id = self._first_state_string(state, self.config.session_id_state_keys)
        if session_id:
            context["session_id"] = session_id
        chat_id = self._state_get(state, "chat_id")
        if chat_id is not None and str(chat_id).strip():
            context["chat_id"] = str(chat_id).strip()

        user_prompt = self._state_get(state, self.config.user_prompt_state_key)
        if user_prompt:
            context["user_prompt"] = _safe_text(user_prompt, max_chars=self.config.reasoning_preview_chars)

        manager_reasoning = self._state_get(state, self.config.manager_reasoning_state_key)
        if manager_reasoning:
            context["manager_reasoning"] = _safe_text(
                manager_reasoning, max_chars=self.config.reasoning_preview_chars
            )

        if reasoning:
            safe_reasoning = _safe_text(reasoning, max_chars=self.config.reasoning_preview_chars)
            context["assistant_reasoning"] = safe_reasoning
            # keep compatibility with existing policy conventions
            context["reasoning"] = safe_reasoning

        if extra:
            context.update(_safe_json(extra))
        return context

    def _emit_event(self, event: Dict[str, Any]) -> None:
        payload = dict(event)
        payload.setdefault("timestamp", _utc_now_iso())
        payload.setdefault("event_type", payload.get("type") or "event")

        if self.config.on_event:
            try:
                self.config.on_event(payload)
            except Exception as exc:
                logger.debug("AgentX on_event callback failed: %s", exc)

        if not self.config.persist_events:
            return
        try:
            self.agendex.ingest_events([payload])
        except Exception as exc:
            logger.debug("Failed to persist AgentX event: %s", exc)

    def _bridge_agendex_event(
        self,
        *,
        event: Dict[str, Any],
        trace_id: str,
        span_id: str,
        parent_span_id: Optional[str],
        actor: str,
        task: str,
        action: str,
        context: Dict[str, Any],
    ) -> Dict[str, Any]:
        event_type = str(event.get("type") or "event")
        bridged: Dict[str, Any] = {
            "timestamp": _utc_now_iso(),
            "event_type": event_type,
            "trace_id": trace_id,
            "span_id": span_id,
            "parent_span_id": parent_span_id,
            "actor": actor,
            "agent_id": self.agendex.agent_id,
            "task": event.get("task") or task,
            "action": event.get("action") or action,
            "context": context,
        }
        bridged.update(_safe_json(event))
        if "decision" not in bridged and "outcome" in bridged:
            bridged["decision"] = bridged.get("outcome")
        if event_type == "invoke_result" and "result" in bridged:
            bridged["result_preview"] = _safe_text(bridged.get("result"), max_chars=800)
        return bridged

    def on_manager_decision(
        self,
        *,
        state: MutableMapping[str, Any],
        llm_input_messages: Optional[list[Any]] = None,
        decision: Any = None,
    ) -> None:
        """Capture manager routing decision as a trace event."""
        trace_id = self._ensure_trace_id(state)
        parent_span_id = self._state_get(state, self._LAST_SPAN_KEY)
        span_id = self._new_span_id()
        reasoning = getattr(decision, "reasoning", None) or self._state_get(
            state, self.config.manager_reasoning_state_key
        )
        routing = self._state_get(state, self.config.manager_decision_state_key) or getattr(
            decision, "next_assistant", None
        )
        task = self._resolve_task(state)
        context = self._build_context(
            state=state,
            reasoning=reasoning,
            extra={"manager_routing_decision": routing},
        )

        event: Dict[str, Any] = {
            "timestamp": _utc_now_iso(),
            "event_type": "manager_decision",
            "trace_id": trace_id,
            "span_id": span_id,
            "parent_span_id": parent_span_id,
            "actor": "manager",
            "agent_id": self.agendex.agent_id,
            "task": task,
            "action": "manager.route",
            "reason": _safe_text(reasoning, max_chars=self.config.reasoning_preview_chars),
            "manager_routing_decision": routing,
            "next_assistant": getattr(decision, "next_assistant", None),
            "required_tool_names": _safe_json(getattr(decision, "required_tool_names", None) or []),
            "context": context,
        }
        message_preview = self._message_preview(llm_input_messages)
        if message_preview:
            event["llm_input_preview"] = message_preview

        self._emit_event(event)
        self._state_set(state, self._MANAGER_SPAN_KEY, span_id)
        self._state_set(state, self._LAST_SPAN_KEY, span_id)

    def on_assistant_llm_response(
        self,
        *,
        assistant_name: str,
        state: MutableMapping[str, Any],
        llm_input_messages: Optional[list[Any]] = None,
        response: Any = None,
    ) -> None:
        """Capture assistant LLM output (including planned tool calls)."""
        trace_id = self._ensure_trace_id(state)
        parent_span_id = self._state_get(state, self._MANAGER_SPAN_KEY) or self._state_get(
            state, self._LAST_SPAN_KEY
        )
        span_id = self._new_span_id()
        task = self._resolve_task(state)
        response_content = getattr(response, "content", response)
        reasoning = _safe_text(response_content, max_chars=self.config.reasoning_preview_chars)
        tool_calls = getattr(response, "tool_calls", None)
        tool_names: list[str] = []
        if isinstance(tool_calls, list):
            for tool_call in tool_calls:
                if isinstance(tool_call, Mapping):
                    name = tool_call.get("name")
                else:
                    name = getattr(tool_call, "name", None)
                if name:
                    tool_names.append(str(name))

        context = self._build_context(
            state=state,
            reasoning=reasoning,
            extra={"assistant_name": assistant_name},
        )
        event: Dict[str, Any] = {
            "timestamp": _utc_now_iso(),
            "event_type": "assistant_llm_result",
            "trace_id": trace_id,
            "span_id": span_id,
            "parent_span_id": parent_span_id,
            "actor": f"assistant:{assistant_name}",
            "agent_id": self.agendex.agent_id,
            "task": task,
            "action": f"assistant.{assistant_name}.llm",
            "reason": reasoning,
            "tool_call_count": len(tool_names),
            "tool_calls": tool_names,
            "context": context,
        }
        message_preview = self._message_preview(llm_input_messages)
        if message_preview:
            event["llm_input_preview"] = message_preview

        self._emit_event(event)
        self._state_set(state, self._assistant_span_key(assistant_name), span_id)
        self._state_set(state, self._LAST_SPAN_KEY, span_id)

    def on_assistant_tool_outputs(
        self,
        *,
        assistant_name: str,
        state: MutableMapping[str, Any],
        tool_outputs: list[Any],
    ) -> None:
        """Capture assistant tool outputs summary."""
        trace_id = self._ensure_trace_id(state)
        parent_span_id = self._state_get(state, self._assistant_span_key(assistant_name)) or self._state_get(
            state, self._LAST_SPAN_KEY
        )
        span_id = self._new_span_id()
        task = self._resolve_task(state)
        preview = [_safe_text(output, max_chars=300) for output in (tool_outputs or [])[:5]]
        context = self._build_context(
            state=state,
            extra={"assistant_name": assistant_name},
        )
        self._emit_event(
            {
                "timestamp": _utc_now_iso(),
                "event_type": "assistant_tool_outputs",
                "trace_id": trace_id,
                "span_id": span_id,
                "parent_span_id": parent_span_id,
                "actor": f"assistant:{assistant_name}",
                "agent_id": self.agendex.agent_id,
                "task": task,
                "action": f"assistant.{assistant_name}.tools",
                "output_count": len(tool_outputs or []),
                "output_preview": preview,
                "context": context,
            }
        )
        self._state_set(state, self._LAST_SPAN_KEY, span_id)

    def wrap_tool(
        self,
        *,
        tool_name: str,
        tool_func: Callable[..., Any],
        state_getter: Optional[Callable[[], Dict[str, Any]]] = None,
    ) -> Callable[..., Any]:
        """Wrap a tool function with Agendex governance and trace correlation."""
        action = self.config.action_map.get(tool_name) or f"{self.config.action_prefix}{tool_name}"

        @wraps(tool_func)
        def wrapped(*args: Any, **kwargs: Any) -> Any:
            total_start = perf_counter()
            state = state_getter() if state_getter else {}
            if not isinstance(state, MutableMapping):
                state = {}

            trace_id = self._ensure_trace_id(state)
            assistant_name = self._state_get(state, self.config.assistant_state_key) or "assistant"
            parent_span_id = self._state_get(state, self._assistant_span_key(str(assistant_name))) or self._state_get(
                state, self._MANAGER_SPAN_KEY
            ) or self._state_get(state, self._LAST_SPAN_KEY)
            span_id = self._new_span_id()
            task = self._resolve_task(state)
            actor = f"tool:{tool_name}"
            params = _safe_json({"args": list(args), "kwargs": kwargs})
            reasoning = self._extract_assistant_reasoning(state)
            context = self._build_context(
                state=state,
                reasoning=reasoning,
                extra={
                    "assistant_name": assistant_name,
                    "tool_name": tool_name,
                },
            )

            def _forward(event: Dict[str, Any]) -> None:
                bridged = self._bridge_agendex_event(
                    event=event,
                    trace_id=trace_id,
                    span_id=span_id,
                    parent_span_id=parent_span_id,
                    actor=actor,
                    task=task,
                    action=action,
                    context=context,
                )
                self._emit_event(bridged)

            governance_start = perf_counter()
            try:
                result = self.agendex.invoke(
                    action=action,
                    params=params if isinstance(params, dict) else {"kwargs": kwargs},
                    task=task,
                    context=context,
                    trace_id=trace_id,
                    span_id=span_id,
                    parent_span_id=parent_span_id,
                    actor=actor,
                    on_event=_forward,
                )
            except (DeniedError, PendingApprovalError):
                self._state_set(state, self._LAST_SPAN_KEY, span_id)
                raise
            except AgendexError as exc:
                governance_latency_ms = (perf_counter() - governance_start) * 1000.0
                if self.config.fallback_on_transport_error:
                    logger.warning("Agendex invoke failed for %s, executing locally: %s", action, exc)
                    local_start = perf_counter()
                    local = tool_func(*args, **kwargs)
                    local_latency_ms = (perf_counter() - local_start) * 1000.0
                    total_latency_ms = (perf_counter() - total_start) * 1000.0
                    self._emit_event(
                        {
                            "timestamp": _utc_now_iso(),
                            "event_type": "tool_result_local_fallback",
                            "trace_id": trace_id,
                            "span_id": span_id,
                            "parent_span_id": parent_span_id,
                            "actor": actor,
                            "agent_id": self.agendex.agent_id,
                            "task": task,
                            "action": action,
                            "reason": str(exc),
                            "latency_ms": local_latency_ms,
                            "governance_latency_ms": governance_latency_ms,
                            "total_latency_ms": total_latency_ms,
                            "result_preview": _safe_text(local, max_chars=800),
                            "context": context,
                        }
                    )
                    self._state_set(state, self._LAST_SPAN_KEY, span_id)
                    return local
                self._state_set(state, self._LAST_SPAN_KEY, span_id)
                raise

            # Shadow mode uses local execution for the actual tool behavior.
            if isinstance(result, Mapping) and result.get("_agendex") == "shadow":
                governance_latency_ms = (perf_counter() - governance_start) * 1000.0
                local_start = perf_counter()
                local_result = tool_func(*args, **kwargs)
                local_latency_ms = (perf_counter() - local_start) * 1000.0
                total_latency_ms = (perf_counter() - total_start) * 1000.0
                self._emit_event(
                    {
                        "timestamp": _utc_now_iso(),
                        "event_type": "tool_result_local",
                        "trace_id": trace_id,
                        "span_id": span_id,
                        "parent_span_id": parent_span_id,
                        "actor": actor,
                        "agent_id": self.agendex.agent_id,
                        "task": task,
                        "action": action,
                        "mode": "shadow",
                        "latency_ms": local_latency_ms,
                        "governance_latency_ms": governance_latency_ms,
                        "total_latency_ms": total_latency_ms,
                        "result_preview": _safe_text(local_result, max_chars=800),
                        "context": context,
                    }
                )
                self._state_set(state, self._LAST_SPAN_KEY, span_id)
                return local_result

            self._state_set(state, self._LAST_SPAN_KEY, span_id)
            return result

        return wrapped


def enable_agentx_governance(
    agendex: AgendexClient,
    *,
    config: Optional[AgentXGovernanceConfig] = None,
) -> AgentXGovernanceHooks:
    """Create governance hooks for AgentX integration."""
    return AgentXGovernanceHooks(agendex=agendex, config=config)


def _import_optional_class(module_name: str, class_name: str) -> Optional[Type[Any]]:
    try:
        module = importlib.import_module(module_name)
        cls = getattr(module, class_name, None)
        if isinstance(cls, type):
            return cls
    except Exception as exc:
        logger.debug("Optional AgentX import failed (%s.%s): %s", module_name, class_name, exc)
    return None


def _active_state_getter() -> Dict[str, Any]:
    state = _ACTIVE_AGENTX_STATE.get()
    if isinstance(state, MutableMapping):
        return state
    return {}


def _patch_init_accept_governance(cls: Type[Any], *, attr_name: str) -> None:
    if getattr(cls, "_agendex_governance_init_patched", False):
        return
    original_init = getattr(cls, "__init__", None)
    if not callable(original_init):
        return

    try:
        init_sig = inspect.signature(original_init)
    except Exception:
        init_sig = None
    accepts_governance = bool(init_sig and "governance_hooks" in init_sig.parameters)

    @wraps(original_init)
    def patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
        local_hooks = kwargs.pop("governance_hooks", None)
        call_kwargs = dict(kwargs)
        if accepts_governance:
            call_kwargs["governance_hooks"] = local_hooks
        original_init(self, *args, **call_kwargs)

        resolved_hooks = local_hooks or getattr(self, attr_name, None) or _DEFAULT_AGENTX_HOOKS
        if resolved_hooks is not None:
            setattr(self, attr_name, resolved_hooks)

        # Best-effort propagation for older AgentBuilder implementations.
        if cls.__name__ == "AgentBuilder" and resolved_hooks is not None:
            manager = getattr(self, "manager", None)
            if manager is not None and getattr(manager, "_governance_hooks", None) is None:
                setattr(manager, "_governance_hooks", resolved_hooks)
            for assistant in getattr(self, "assistants", []) or []:
                if getattr(assistant, "_governance_hooks", None) is None:
                    setattr(assistant, "_governance_hooks", resolved_hooks)
            attach = getattr(self, "_attach_governance_hooks", None)
            if callable(attach):
                try:
                    attach()
                except Exception as exc:
                    logger.debug("AgentBuilder governance attach failed: %s", exc)

    cls.__init__ = patched_init
    setattr(cls, "_agendex_governance_init_patched", True)


def _patch_tool_builder(tool_builder_cls: Type[Any]) -> None:
    if getattr(tool_builder_cls, "_agendex_governance_tool_patched", False):
        return

    original_init = getattr(tool_builder_cls, "__init__", None)
    original_wrap = getattr(tool_builder_cls, "_wrap_with_context", None)
    if not callable(original_init) or not callable(original_wrap):
        return

    try:
        init_sig = inspect.signature(original_init)
    except Exception:
        init_sig = None
    accepts_governance = bool(init_sig and "governance_hooks" in init_sig.parameters)
    accepts_state_getter = bool(init_sig and "state_getter" in init_sig.parameters)

    @wraps(original_init)
    def patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
        local_hooks = kwargs.pop("governance_hooks", None)
        local_state_getter = kwargs.pop("state_getter", None)

        call_kwargs = dict(kwargs)
        if accepts_governance:
            call_kwargs["governance_hooks"] = local_hooks
        if accepts_state_getter:
            call_kwargs["state_getter"] = local_state_getter
        original_init(self, *args, **call_kwargs)

        if local_hooks is None:
            local_hooks = getattr(self, "governance_hooks", None)
        if local_hooks is None:
            local_hooks = _DEFAULT_AGENTX_HOOKS
        setattr(self, "governance_hooks", local_hooks)
        setattr(self, "state_getter", local_state_getter or getattr(self, "state_getter", None))

    tool_builder_cls.__init__ = patched_init

    if not _method_references_name(original_wrap, "wrap_tool"):

        @wraps(original_wrap)
        def patched_wrap(self: Any, func: Callable[..., Any]) -> Callable[..., Any]:
            wrapped = original_wrap(self, func)
            hooks = getattr(self, "governance_hooks", None) or _DEFAULT_AGENTX_HOOKS
            if not hooks or not hasattr(hooks, "wrap_tool"):
                return wrapped
            state_getter = getattr(self, "state_getter", None) or _active_state_getter
            tool_name = getattr(func, "__name__", getattr(wrapped, "__name__", "tool"))
            try:
                return hooks.wrap_tool(
                    tool_name=str(tool_name),
                    tool_func=wrapped,
                    state_getter=state_getter,
                )
            except Exception as exc:
                logger.debug("ToolBuilder governance shim failed for %s: %s", tool_name, exc)
                return wrapped

        tool_builder_cls._wrap_with_context = patched_wrap

    setattr(tool_builder_cls, "_agendex_governance_tool_patched", True)


def _patch_run_node_with_state_context(
    cls: Type[Any],
    *,
    emit_manager_event: bool = False,
    emit_assistant_events: bool = False,
) -> None:
    marker = "_agendex_governance_run_node_patched"
    if getattr(cls, marker, False):
        return
    original_run_node = getattr(cls, "run_node", None)
    if not callable(original_run_node):
        return

    has_manager_emit = _method_references_name(original_run_node, "on_manager_decision")
    has_assistant_llm_emit = _method_references_name(original_run_node, "on_assistant_llm_response")
    has_assistant_tools_emit = _method_references_name(original_run_node, "on_assistant_tool_outputs")

    @wraps(original_run_node)
    def patched_run_node(self: Any, state: Any, *args: Any, **kwargs: Any) -> Any:
        token = _ACTIVE_AGENTX_STATE.set(state)
        try:
            result = original_run_node(self, state, *args, **kwargs)
            hooks = (
                getattr(self, "_governance_hooks", None)
                or getattr(self, "governance_hooks", None)
                or _DEFAULT_AGENTX_HOOKS
            )
            if hooks and emit_manager_event and not has_manager_emit and hasattr(hooks, "on_manager_decision"):
                try:
                    hooks.on_manager_decision(
                        state=state,
                        llm_input_messages=state.get("manager_node_context_messages")
                        if isinstance(state, MutableMapping)
                        else None,
                        decision=None,
                    )
                except Exception as exc:
                    logger.debug("Manager fallback governance emit failed: %s", exc)

            if hooks and emit_assistant_events:
                assistant_name = getattr(self, "_assistant_name", None) or getattr(self, "name", "assistant")
                if not has_assistant_llm_emit and hasattr(hooks, "on_assistant_llm_response"):
                    try:
                        hooks.on_assistant_llm_response(
                            assistant_name=str(assistant_name),
                            state=state,
                            llm_input_messages=state.get("assistant_messages_current_turn")
                            if isinstance(state, MutableMapping)
                            else None,
                            response=_extract_assistant_response_from_state(state),
                        )
                    except Exception as exc:
                        logger.debug("Assistant fallback LLM governance emit failed: %s", exc)
                if not has_assistant_tools_emit and hasattr(hooks, "on_assistant_tool_outputs"):
                    outputs = _extract_tool_outputs_from_state(state)
                    if outputs:
                        try:
                            hooks.on_assistant_tool_outputs(
                                assistant_name=str(assistant_name),
                                state=state,
                                tool_outputs=outputs,
                            )
                        except Exception as exc:
                            logger.debug("Assistant fallback tool-output governance emit failed: %s", exc)
            return result
        finally:
            _ACTIVE_AGENTX_STATE.reset(token)

    cls.run_node = patched_run_node
    setattr(cls, marker, True)


def install_agentx_runtime_shims(
    hooks: AgentXGovernanceHooks,
    *,
    agent_builder_cls: Optional[Type[Any]] = None,
) -> None:
    """
    Install compatibility shims for AgentX so governance can be enabled with
    minimal client changes, including older builds without explicit hook params.
    """
    global _DEFAULT_AGENTX_HOOKS
    _DEFAULT_AGENTX_HOOKS = hooks

    tool_builder_cls = _import_optional_class("agentx.core.tools.tool_builder", "ToolBuilder")
    if tool_builder_cls:
        _patch_tool_builder(tool_builder_cls)

    manager_builder_cls = _import_optional_class(
        "agentx.core.graph.nodes.dynamic.manager.manager_builder", "ManagerBuilder"
    )
    if manager_builder_cls:
        _patch_init_accept_governance(manager_builder_cls, attr_name="_governance_hooks")
        _patch_run_node_with_state_context(manager_builder_cls, emit_manager_event=True)

    assistant_builder_cls = _import_optional_class(
        "agentx.core.graph.nodes.dynamic.assistants.assistant_builder", "AssistantBuilder"
    )
    if assistant_builder_cls:
        _patch_init_accept_governance(assistant_builder_cls, attr_name="_governance_hooks")
        _patch_run_node_with_state_context(assistant_builder_cls, emit_assistant_events=True)

    manager_tool_executor_cls = _import_optional_class(
        "agentx.core.graph.nodes.static.manager_tool_executor", "ManagerToolExecutorNode"
    )
    if manager_tool_executor_cls:
        _patch_init_accept_governance(manager_tool_executor_cls, attr_name="_governance_hooks")
        _patch_run_node_with_state_context(manager_tool_executor_cls)

    resolved_agent_builder_cls = agent_builder_cls or _import_optional_class(
        "agentx.core.graph.agent_builder", "AgentBuilder"
    )
    if resolved_agent_builder_cls:
        _patch_init_accept_governance(resolved_agent_builder_cls, attr_name="governance_hooks")


class GovernedAgentBuilder:
    """
    Drop-in wrapper around AgentX AgentBuilder that auto-wires Agendex hooks.

    Client-side usage:
        builder = GovernedAgentBuilder(...same AgentBuilder args...)
        graph = builder.build_agent()
    """

    def __init__(
        self,
        *agent_builder_args: Any,
        agendex: Optional[AgendexClient] = None,
        agendex_client_kwargs: Optional[Dict[str, Any]] = None,
        governance_config: Optional[AgentXGovernanceConfig] = None,
        agent_builder_cls: Optional[Type[Any]] = None,
        install_runtime_shims_flag: bool = True,
        **agent_builder_kwargs: Any,
    ) -> None:
        self.agendex = agendex or AgendexClient(**(agendex_client_kwargs or {}))
        self.hooks = enable_agentx_governance(self.agendex, config=governance_config)

        if install_runtime_shims_flag:
            install_agentx_runtime_shims(self.hooks, agent_builder_cls=agent_builder_cls)

        resolved_cls = agent_builder_cls or _import_optional_class("agentx.core.graph.agent_builder", "AgentBuilder")
        if not resolved_cls:
            raise ImportError(
                "Could not import AgentX AgentBuilder. Pass `agent_builder_cls=` explicitly "
                "or ensure agentx is installed in the runtime."
            )
        self._agent_builder_cls = resolved_cls
        self._builder = self._build_underlying(agent_builder_args, agent_builder_kwargs)

    def _build_underlying(self, args: tuple[Any, ...], kwargs: Dict[str, Any]) -> Any:
        try:
            call_kwargs = dict(kwargs)
            call_kwargs.setdefault("governance_hooks", self.hooks)
            return self._agent_builder_cls(*args, **call_kwargs)
        except TypeError as exc:
            if "governance_hooks" not in str(exc):
                raise
            builder = self._agent_builder_cls(*args, **kwargs)
            self._attach_hooks_to_legacy_builder(builder)
            return builder

    def _attach_hooks_to_legacy_builder(self, builder: Any) -> None:
        setattr(builder, "governance_hooks", self.hooks)
        manager = getattr(builder, "manager", None)
        if manager is not None and getattr(manager, "_governance_hooks", None) is None:
            setattr(manager, "_governance_hooks", self.hooks)
        for assistant in getattr(builder, "assistants", []) or []:
            if getattr(assistant, "_governance_hooks", None) is None:
                setattr(assistant, "_governance_hooks", self.hooks)

    @property
    def builder(self) -> Any:
        """Access the underlying AgentBuilder instance."""
        return self._builder

    def build_agent(self) -> Any:
        return self._builder.build_agent()

    def __getattr__(self, item: str) -> Any:
        return getattr(self._builder, item)


__all__ = [
    "AgentXGovernanceConfig",
    "AgentXGovernanceHooks",
    "GovernedAgentBuilder",
    "enable_agentx_governance",
    "install_agentx_runtime_shims",
]
