"""
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 logging
from dataclasses import dataclass, field
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
from .decision_capture import (
    AgendexCallbackHandler,
    DecisionAccumulator,
    DecisionContextBuilder,
    build_transport_status,
    extract_message_content,
    normalize_messages,
    safe_json,
    safe_text,
    utc_now_iso,
)

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 _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)
    # Maps tool_name -> (semantic_action, params_extractor(*args, **kwargs) -> dict)
    tool_resource_map: Dict[str, Any] = 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"
    user_prompt_fallback_keys: tuple[str, ...] = ("user_prompt", "prompt", "input", "query", "user_input")
    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"
    assistant_messages_fallback_keys: tuple[str, ...] = ("messages", "assistant_messages", "history", "chat_history")
    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")
    decision_accumulator_state_key: str = "_agendex_decision_accumulator"

    # Behavior
    persist_events: bool = True
    capture_message_preview: bool = False
    capture_extra_state: bool = True
    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()
        self._decision_builder = DecisionContextBuilder(
            user_id_keys=self.config.user_id_state_keys,
            session_id_keys=self.config.session_id_state_keys,
            user_intent_keys=(self.config.user_prompt_state_key, *self.config.user_prompt_fallback_keys),
            assistant_rationale_keys=("assistant_rationale", "assistant_reasoning", "reasoning", self.config.manager_reasoning_state_key),
            routing_choice_keys=("routing_choice", self.config.manager_decision_state_key, "next_assistant", self.config.assistant_state_key),
            assistant_name_keys=("assistant_name", self.config.assistant_state_key),
            extra_state_max_chars=self.config.reasoning_preview_chars,
            message_preview_chars=self.config.message_preview_chars,
        )

    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 _get_accumulator(self, state: Any) -> DecisionAccumulator:
        accumulator = self._state_get(state, self.config.decision_accumulator_state_key)
        if isinstance(accumulator, DecisionAccumulator):
            return accumulator
        accumulator = DecisionAccumulator()
        self._state_set(state, self.config.decision_accumulator_state_key, accumulator)
        return accumulator

    def create_callback_handler(self, state: MutableMapping[str, Any]) -> AgendexCallbackHandler:
        return AgendexCallbackHandler(self._get_accumulator(state))

    def _extract_messages(self, state: Any) -> Optional[list[Any]]:
        messages = self._state_get(state, self.config.assistant_messages_state_key)
        if isinstance(messages, list) and messages:
            return messages
        for key in self.config.assistant_messages_fallback_keys:
            messages = self._state_get(state, key)
            if isinstance(messages, list) and messages:
                return messages
        return None

    def _extract_assistant_reasoning(self, state: Any) -> Optional[str]:
        messages = self._extract_messages(state)
        if not messages:
            return None
        for msg in reversed(messages):
            content = extract_message_content(msg)
            if content is None:
                continue
            text = safe_text(content, max_chars=self.config.reasoning_preview_chars)
            if text:
                return text
        return None

    def _message_preview(self, messages: Any) -> Optional[list[str]]:
        if not self.config.capture_message_preview:
            return None
        return normalize_messages(messages, max_chars=self.config.message_preview_chars)

    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 not user_prompt:
            for key in self.config.user_prompt_fallback_keys:
                user_prompt = self._state_get(state, key)
                if user_prompt:
                    break
        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))

        if self.config.capture_extra_state and isinstance(state, MutableMapping):
            known_keys = {
                self.config.user_prompt_state_key,
                self.config.manager_reasoning_state_key,
                self.config.manager_decision_state_key,
                self.config.assistant_state_key,
                self.config.assistant_messages_state_key,
                self.config.task_state_key,
                self.config.trace_state_key,
                self.config.trace_context_key,
                self.config.decision_accumulator_state_key,
                *self.config.user_id_state_keys,
                *self.config.session_id_state_keys,
                *self.config.user_prompt_fallback_keys,
                *self.config.assistant_messages_fallback_keys,
            }
            for key, value in state.items():
                if key in context or key in known_keys or key.startswith("_agendex"):
                    continue
                if isinstance(value, (str, int, float, bool)):
                    context.setdefault(key, value)
                elif isinstance(value, Mapping):
                    context.setdefault(key, safe_json(value))
        return context

    def _build_capture(
        self,
        *,
        decision_type: str,
        state: Any,
        action: str,
        task: str,
        trace_id: str,
        span_id: str,
        parent_span_id: Optional[str],
        actor: str,
        params: Optional[Dict[str, Any]] = None,
        assistant_name: Optional[str] = None,
        tool_name: Optional[str] = None,
        llm_input_messages: Any = None,
        assistant_response: Any = None,
        manager_output: Any = None,
        tool_call_payload: Any = None,
        tool_action_intent: Optional[str] = None,
        transport_status: Optional[Dict[str, Any]] = None,
    ):
        accumulator = self._get_accumulator(state)
        context = self._build_context(
            state=state,
            reasoning=self._extract_assistant_reasoning(state),
            extra={
                "assistant_name": assistant_name,
                "tool_name": tool_name,
                "trace_id": trace_id,
                "span_id": span_id,
                "parent_span_id": parent_span_id,
                "actor": actor,
                "mode": getattr(self.agendex, "mode", None),
            },
        )
        return self._decision_builder.build(
            decision_type=decision_type,
            action=action,
            task=task,
            context=context,
            state=state,
            params=params or {},
            accumulator=accumulator,
            trace_id=trace_id,
            span_id=span_id,
            parent_span_id=parent_span_id,
            actor=actor,
            agent_id=self.agendex.agent_id,
            tenant_id=getattr(self.agendex, "tenant_id", None),
            mode=getattr(self.agendex, "mode", None),
            assistant_name=assistant_name,
            tool_name=tool_name,
            messages=llm_input_messages,
            assistant_response=assistant_response,
            manager_output=manager_output,
            tool_call_payload=tool_call_payload,
            tool_action_intent=tool_action_intent,
            transport_status=transport_status,
        )

    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 _flush_events(self, events: list[Dict[str, Any]]) -> None:
        for event in events:
            self._emit_event(event)

    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)
        if reasoning is None and isinstance(decision, Mapping):
            reasoning = decision.get("reasoning")
        if not reasoning:
            reasoning = self._state_get(state, self.config.manager_reasoning_state_key)
        routing = self._state_get(state, self.config.manager_decision_state_key)
        if not routing:
            routing = getattr(decision, "next_assistant", None)
            if routing is None and isinstance(decision, Mapping):
                routing = decision.get("next_assistant")
        task = self._resolve_task(state)
        required_tool_names = (
            getattr(decision, "required_tool_names", None)
            or (decision.get("required_tool_names") if isinstance(decision, Mapping) else None)
            or []
        )
        accumulator = self._get_accumulator(state)
        accumulator.record_prompt(messages=llm_input_messages)
        accumulator.record_manager_decision(
            reasoning=reasoning,
            routing_choice=routing,
            required_tool_names=required_tool_names,
        )
        capture = self._build_capture(
            decision_type="manager_route",
            state=state,
            action="manager.route",
            task=task,
            trace_id=trace_id,
            span_id=span_id,
            parent_span_id=parent_span_id,
            actor="manager",
            assistant_name=str(routing) if routing else None,
            llm_input_messages=llm_input_messages,
            manager_output={
                "reasoning": reasoning,
                "routing_choice": routing,
                "required_tool_names": required_tool_names,
            },
            tool_action_intent=None,
            transport_status=build_transport_status(),
        )
        context = capture.as_context()

        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)
            if not isinstance(decision, Mapping)
            else decision.get("next_assistant"),
            "required_tool_names": safe_json(required_tool_names),
            "context": context,
            "transport_status": build_transport_status(),
        }
        event.update(capture.as_event_fields())
        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 = extract_message_content(response) if response is not None else None
        if response_content is None:
            response_content = response
        reasoning = safe_text(response_content, max_chars=self.config.reasoning_preview_chars)
        tool_calls = getattr(response, "tool_calls", None)
        if tool_calls is None and isinstance(response, Mapping):
            tool_calls = response.get("tool_calls")
        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))

        accumulator = self._get_accumulator(state)
        accumulator.record_prompt(messages=llm_input_messages)
        accumulator.record_response(
            response=response,
            assistant_name=assistant_name,
            tool_calls=tool_calls,
        )
        capture = self._build_capture(
            decision_type="tool_call" if tool_names else "assistant_llm_result",
            state=state,
            action=f"assistant.{assistant_name}.llm",
            task=task,
            trace_id=trace_id,
            span_id=span_id,
            parent_span_id=parent_span_id,
            actor=f"assistant:{assistant_name}",
            assistant_name=assistant_name,
            llm_input_messages=llm_input_messages,
            assistant_response=response,
            tool_call_payload=tool_calls,
            tool_action_intent=tool_names[0] if tool_names else None,
            transport_status=build_transport_status(),
        )
        context = capture.as_context()
        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,
            "transport_status": build_transport_status(),
        }
        event.update(capture.as_event_fields())
        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]]
        accumulator = self._get_accumulator(state)
        accumulator.merge(
            assistant_name=assistant_name,
            assistant_response_excerpt=preview[-1] if preview else None,
            recent_messages=tool_outputs[-5:] if tool_outputs else None,
        )
        capture = self._build_capture(
            decision_type="assistant_tool_outputs",
            state=state,
            action=f"assistant.{assistant_name}.tools",
            task=task,
            trace_id=trace_id,
            span_id=span_id,
            parent_span_id=parent_span_id,
            actor=f"assistant:{assistant_name}",
            assistant_name=assistant_name,
            tool_call_payload={"output_preview": preview},
            transport_status=build_transport_status(),
        )
        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": capture.as_context(),
                "transport_status": build_transport_status(),
                **capture.as_event_fields(),
            }
        )
        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."""
        if tool_name in self.config.tool_resource_map:
            semantic_action, params_extractor = self.config.tool_resource_map[tool_name]
            action = semantic_action
        else:
            params_extractor = None
            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}"
            if params_extractor is not None:
                params = safe_json(params_extractor(*args, **kwargs))
            else:
                params = safe_json({"args": list(args), "kwargs": kwargs})
            params_payload = params if isinstance(params, dict) else {"kwargs": kwargs}
            accumulator = self._get_accumulator(state)
            accumulator.merge(
                assistant_name=assistant_name,
                tool_name=tool_name,
                tool_call_payload={"action": action, "params": safe_json(params_payload)},
                timestamps={"last_tool_boundary_at": utc_now_iso()},
            )
            capture = self._build_capture(
                decision_type="tool_call",
                state=state,
                action=action,
                task=task,
                trace_id=trace_id,
                span_id=span_id,
                parent_span_id=parent_span_id,
                actor=actor,
                params=params_payload,
                assistant_name=str(assistant_name),
                tool_name=tool_name,
                llm_input_messages=self._extract_messages(state),
                assistant_response=self._extract_assistant_reasoning(state),
                manager_output={
                    "reasoning": self._state_get(state, self.config.manager_reasoning_state_key),
                    "routing_choice": self._state_get(state, self.config.manager_decision_state_key),
                },
                tool_call_payload={"action": action, "params": safe_json(params_payload)},
                transport_status=build_transport_status(),
            )
            context = capture.as_context()
            forwarded_events: list[Dict[str, Any]] = []

            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,
                )
                forwarded_events.append(bridged)

            governance_start = perf_counter()
            try:
                result = self.agendex.invoke(
                    action=action,
                    params=params_payload,
                    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._flush_events(forwarded_events)
                self._state_set(state, self._LAST_SPAN_KEY, span_id)
                raise
            except AgendexError as exc:
                governance_round_trip_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
                    transport_status = build_transport_status(ok=False, error=str(exc), fallback=True)
                    self._flush_events(forwarded_events)
                    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,
                            "policy_latency_ms": None,
                            "governance_round_trip_ms": governance_round_trip_ms,
                            "tool_latency_ms": local_latency_ms,
                            "total_latency_ms": total_latency_ms,
                            "result_preview": safe_text(local, max_chars=800),
                            "context": context,
                            "transport_status": transport_status,
                            **capture.as_event_fields(),
                        }
                    )
                    self._state_set(state, self._LAST_SPAN_KEY, span_id)
                    return local
                self._flush_events(forwarded_events)
                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":
                local_start = perf_counter()
                local_result = tool_func(*args, **kwargs)
                local_latency_ms = (perf_counter() - local_start) * 1000.0
                policy_latency_ms = None
                governance_round_trip_ms = None
                transport_status = build_transport_status()
                shadow_future = result.get("_shadow_future")
                if shadow_future is not None:
                    try:
                        bg_result = shadow_future.result(timeout=5.0)
                        if isinstance(bg_result, Mapping):
                            policy_latency_ms = bg_result.get("server_latency_ms")
                            governance_round_trip_ms = bg_result.get("round_trip_ms")
                            returned_transport = bg_result.get("transport_status")
                            if isinstance(returned_transport, Mapping):
                                transport_status = dict(returned_transport)
                    except Exception:
                        transport_status = build_transport_status(ok=False, error="shadow future collection failed", fallback=True)
                total_latency_ms = (perf_counter() - total_start) * 1000.0
                self._flush_events(forwarded_events)
                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,
                        "policy_latency_ms": policy_latency_ms,
                        "governance_round_trip_ms": governance_round_trip_ms,
                        "tool_latency_ms": local_latency_ms,
                        "total_latency_ms": total_latency_ms,
                        "result_preview": safe_text(local_result, max_chars=800),
                        "context": context,
                        "transport_status": transport_status,
                        **capture.as_event_fields(),
                    }
                )
                self._state_set(state, self._LAST_SPAN_KEY, span_id)
                return local_result

            self._flush_events(forwarded_events)
            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",
]
