import asyncio
import functools
import hashlib
import inspect
import json
import os
import uuid
from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple

import requests

from .sdk_adk_instance_patch import patch_adk_agent_instance

from .runtime import kyb_client, resolve_run_id_from_invocation_id
from .context import kyb_run_id
from ._asyncutil import fire_and_forget
from .auto import install as install_http_patches
from .adapters import wrap_adk_tool
from .spans import span as kyb_span
from .intent_fingerprint import build_intent_key

try:
    from google.adk.tools.base_tool import BaseTool  # type: ignore
except Exception:
    BaseTool = None  # type: ignore

def _is_adk_basetool(obj: Any) -> bool:
        return BaseTool is not None and isinstance(obj, BaseTool)

def _mark(obj: Any, name: str) -> bool:
    """Returns True if already marked, else marks and returns False."""
    if getattr(obj, name, False):
        return True
    try:
        setattr(obj, name, True)
    except Exception:
        object.__setattr__(obj, name, True)
    return False

class KybernisProtocolException(RuntimeError):
    pass


class KybernisSDK:
    _DEFAULT_SEMANTIC_PATHS: Tuple[str, ...] = ("tool_name", "tool_args", "session_signature")
    _DEFAULT_NOISY_SESSION_KEYS: Set[str] = {
        "timestamp",
        "ts",
        "created_at",
        "updated_at",
        "last_updated",
        "last_seen",
        "request_id",
        "trace_id",
        "span_id",
        "event_id",
        "message_id",
        "nonce",
        "latency_ms",
        "duration_ms",
        "prompt_tokens",
        "completion_tokens",
        "total_tokens",
        "usage",
        "usage_tokens",
    }
    _MANDATE_PRESETS: Dict[str, Dict[str, Any]] = {
        "money": {
            "semantic_paths": [
                "tool_name",
                "tool_args.customer_id",
                "tool_args.order_id",
                "tool_args.payment_id",
                "tool_args.amount_cents",
                "tool_args.currency",
                "tool_args.reason",
                "session_signature",
            ],
            "session_signature_mode": "stable_subset",
        },
        "infra": {
            "semantic_paths": [
                "tool_name",
                "tool_args.action",
                "tool_args.resource_id",
                "tool_args.environment",
                "tool_args.change_id",
                "session_signature",
            ],
            "session_signature_mode": "stable_subset",
        },
        "low_risk": {
            "semantic_paths": [
                "tool_name",
                "tool_args",
            ],
            "session_signature_mode": "off",
        },
    }

    def __init__(
        self,
        framework_instance: Any = None,
        api_key: str = "",
        gateway_url: str = "",
        tenant_id: str = "",
        timeout_s: int = 10,
        auto_intercept: bool = False,
        proxy_mode: bool = False,
        mode: str = "",
        allow_methods: Optional[Iterable[str]] = None,
        exclude_methods: Optional[Iterable[str]] = None,
        tool_map: Optional[Dict[str, str]] = None,
        tool_resolver: Optional[Callable[[str, Any, tuple, Dict[str, Any]], Optional[str]]] = None,
        intent_context_provider: Optional[Callable[[str, Dict[str, Any], Any], Dict[str, Any]]] = None,
        drift_guard_enabled: Optional[bool] = None,
        drift_guard_mode: str = "",
        policy_version: str = "",
        dependency_signature: str = "",
        model_label: str = "",
        provider_label: str = "",
        prompt_template_version: str = "",
        tool_schema_version: str = "",
        decoding_params: Optional[Dict[str, Any]] = None,
        behavior_probe_mode: str = "",
        behavior_probe_suite_version: str = "",
        behavior_probe_signature: str = "",
        behavior_probe_checked_at: str = "",
        context_mode: str = "",
        semantic_paths: Optional[Iterable[str]] = None,
        session_signature_mode: str = "",
        mandate_preset: str = "",
    ) -> None:
        if isinstance(framework_instance, str) and not api_key:
            api_key = framework_instance
            framework_instance = None
        self.api_key = (
            api_key
            or os.getenv("KYB_API_KEY", "")
            or os.getenv("KYB_API_TOKEN", "")
            or os.getenv("KYBERNIS_ADMIN_TOKEN", "")
        )
        self.gateway_url = gateway_url.strip() or os.getenv("KYB_PROTOCOL_GATEWAY", "") or os.getenv("KYB_API_BASE", "")
        if self.gateway_url:
            self.gateway_url = self.gateway_url.rstrip("/")
        self.tenant_id = tenant_id.strip() or os.getenv("KYB_TENANT_ID", "")
        self.timeout_s = timeout_s
        env_auto = os.getenv("KYB_SDK_AUTO_INTERCEPT", "").strip().lower()
        self.auto_intercept = auto_intercept or env_auto in {"1", "true", "yes", "on"}
        env_proxy = os.getenv("KYB_SDK_PROXY_MODE", "").strip().lower()
        self.proxy_mode = proxy_mode or env_proxy in {"1", "true", "yes", "on"}
        self.mode_proxy, self.mode_mcp = self._resolve_mode(mode)
        if self.mode_proxy:
            install_http_patches()
        env_allow = os.getenv("KYB_SDK_METHOD_ALLOWLIST", "")
        env_block = os.getenv("KYB_SDK_METHOD_BLOCKLIST", "")
        self.allow_methods: Set[str] = set(allow_methods or [])
        if env_allow:
            self.allow_methods.update(self._split_csv(env_allow))
        self.exclude_methods: Set[str] = set(exclude_methods or [])
        if env_block:
            self.exclude_methods.update(self._split_csv(env_block))
        self.tool_map: Dict[str, str] = dict(tool_map or {})
        env_tool_map = os.getenv("KYB_SDK_METHOD_TOOL_MAP", "").strip()
        if env_tool_map:
            try:
                data = json.loads(env_tool_map)
                if isinstance(data, dict):
                    for key, value in data.items():
                        if isinstance(key, str) and isinstance(value, str):
                            self.tool_map[key] = value
            except Exception:
                pass
        self.tool_resolver = tool_resolver
        self.intent_context_provider = intent_context_provider
        if drift_guard_enabled is None:
            self.drift_guard_enabled = self._env_bool("KYB_DRIFT_GUARD_ENABLED", True)
        else:
            self.drift_guard_enabled = bool(drift_guard_enabled)
        mode_raw = (drift_guard_mode or os.getenv("KYB_DRIFT_GUARD_MODE", "warn_only")).strip().lower()
        self.drift_guard_mode = "soft_fail" if mode_raw == "soft_fail" else "warn_only"
        self.policy_version = (policy_version or os.getenv("KYB_POLICY_VERSION", "")).strip()
        self.dependency_signature = (dependency_signature or os.getenv("KYB_DEPENDENCY_SIGNATURE", "")).strip()
        self.model_label = (model_label or os.getenv("KYB_MODEL_LABEL", "")).strip()
        self.provider_label = (provider_label or os.getenv("KYB_PROVIDER_LABEL", "")).strip()
        self.prompt_template_version = (
            prompt_template_version or os.getenv("KYB_PROMPT_TEMPLATE_VERSION", "")
        ).strip()
        self.tool_schema_version = (tool_schema_version or os.getenv("KYB_TOOL_SCHEMA_VERSION", "")).strip()
        raw_decoding: Any = decoding_params
        if raw_decoding is None:
            raw_decoding = os.getenv("KYB_DECODING_PARAMS_JSON", "")
        self.decoding_params = self._parse_json_object(raw_decoding)
        raw_probe_mode = (behavior_probe_mode or os.getenv("KYB_BEHAVIOR_PROBE_MODE", "off")).strip().lower()
        if raw_probe_mode not in {"off", "observe", "enforce"}:
            raw_probe_mode = "off"
        self.behavior_probe_mode = raw_probe_mode
        self.behavior_probe_suite_version = (
            behavior_probe_suite_version or os.getenv("KYB_BEHAVIOR_PROBE_SUITE_VERSION", "")
        ).strip()
        self.behavior_probe_signature = (
            behavior_probe_signature or os.getenv("KYB_BEHAVIOR_PROBE_SIGNATURE", "")
        ).strip()
        self.behavior_probe_checked_at = (
            behavior_probe_checked_at or os.getenv("KYB_BEHAVIOR_PROBE_CHECKED_AT", "")
        ).strip()
        raw_preset = (mandate_preset or os.getenv("KYB_MANDATE_PRESET", "")).strip().lower()
        self.mandate_preset = self._normalize_mandate_preset(raw_preset)
        preset = self.get_mandate_preset(self.mandate_preset) if self.mandate_preset else {}
        raw_context_mode = (context_mode or os.getenv("KYB_CONTEXT_MODE", "auto")).strip().lower()
        self.context_mode = self._normalize_context_mode(raw_context_mode)
        env_semantic_paths = self._parse_csv_list(os.getenv("KYB_SEMANTIC_PATHS", ""))
        user_semantic_paths = [p.strip() for p in (semantic_paths or []) if isinstance(p, str) and p.strip()]
        merged_paths = list(user_semantic_paths) if user_semantic_paths else list(env_semantic_paths)
        if not merged_paths:
            preset_paths = preset.get("semantic_paths")
            if isinstance(preset_paths, list):
                merged_paths = [str(p).strip() for p in preset_paths if str(p).strip()]
        if not merged_paths:
            merged_paths = list(self._DEFAULT_SEMANTIC_PATHS)
        self.semantic_paths: List[str] = merged_paths
        explicit_sig_mode = (session_signature_mode or os.getenv("KYB_SESSION_SIGNATURE_MODE", "")).strip().lower()
        if explicit_sig_mode:
            raw_sig_mode = explicit_sig_mode
        else:
            raw_sig_mode = str(preset.get("session_signature_mode") or "stable_subset").strip().lower()
        self.session_signature_mode = self._normalize_session_signature_mode(raw_sig_mode)
        if framework_instance is not None:
            self.wrap_framework(framework_instance)

    def wrap_framework(self, framework_instance: Any):
        framework = self._detect_framework(framework_instance)
        if self.proxy_mode:
            return _SDKProxy(self, framework_instance, framework)
        if framework == "langgraph" and hasattr(framework_instance, "invoke"):
            return self._intercept(framework_instance, "invoke", framework=framework)
        if framework == "crewai" and hasattr(framework_instance, "kickoff"):
            return self._intercept(framework_instance, "kickoff", framework=framework)
        if framework == "autogen" and hasattr(framework_instance, "initiate_chat"):
            return self._intercept(framework_instance, "initiate_chat", framework=framework)
        if framework == "google_adk":
            print("Running patches though")
            self._wrap_adk_tools(framework_instance)
            patch_adk_agent_instance(framework_instance, also_patch_subagents=True)
            
            return framework_instance
        if self.auto_intercept:
            return self._intercept_all_methods(framework_instance, framework=framework)
        return framework_instance

    @staticmethod
    def _env_bool(name: str, default: bool) -> bool:
        raw = os.getenv(name, "").strip().lower()
        if raw in {"1", "true", "yes", "on"}:
            return True
        if raw in {"0", "false", "no", "off"}:
            return False
        return default

    @staticmethod
    def _is_ambiguous_execution_error(exc: BaseException) -> bool:
        if isinstance(exc, (TimeoutError, asyncio.TimeoutError)):
            return True
        if isinstance(exc, (requests.exceptions.Timeout, requests.exceptions.ConnectionError)):
            return True
        name = type(exc).__name__.strip().lower()
        if any(tok in name for tok in ("timeout", "connectionerror", "readtimeout", "connecttimeout", "transport")):
            return True
        text = str(exc).strip().lower()
        markers = (
            "timeout",
            "timed out",
            "deadline exceeded",
            "connection reset",
            "connection aborted",
            "connection refused",
            "broken pipe",
            "temporarily unavailable",
            "eof",
            "unknown outcome",
            "socket closed",
        )
        return any(marker in text for marker in markers)

    @staticmethod
    def _unknown_outcome_reason(exc: BaseException) -> str:
        name = type(exc).__name__.strip() or "Error"
        text = str(exc).strip()
        if not text:
            return f"outcome_unknown:{name}"
        if len(text) > 240:
            text = text[:240]
        return f"outcome_unknown:{name}:{text}"

    @staticmethod
    def _normalize_context_mode(raw: str) -> str:
        # Today "auto" and "adk_default" are equivalent because Google ADK is
        # the first built-in orchestrator integration. Keep "auto" as alias
        # so future framework-aware routing can evolve without API breakage.
        if raw in {"auto", ""}:
            return "adk_default"
        if raw in {"provider", "custom"}:
            return "provider"
        if raw in {"adk_default", "default", "builtin"}:
            return "adk_default"
        return "adk_default"

    @staticmethod
    def _normalize_session_signature_mode(raw: str) -> str:
        if raw in {"off", "none", "disabled"}:
            return "off"
        if raw in {"full", "all"}:
            return "full"
        return "stable_subset"

    @classmethod
    def _normalize_mandate_preset(cls, raw: str) -> str:
        name = (raw or "").strip().lower().replace("-", "_")
        if not name:
            return ""
        if name in cls._MANDATE_PRESETS:
            return name
        return ""

    @classmethod
    def get_mandate_preset(cls, name: str) -> Dict[str, Any]:
        """
        Preset helper for default mandate-context configuration.
        Returns:
          {
            "semantic_paths": [...],
            "session_signature_mode": "...",
          }
        """
        normalized = cls._normalize_mandate_preset(name)
        if not normalized:
            return {}
        preset = cls._MANDATE_PRESETS.get(normalized) or {}
        return {
            "semantic_paths": list(preset.get("semantic_paths") or []),
            "session_signature_mode": str(preset.get("session_signature_mode") or "stable_subset"),
        }

    @staticmethod
    def _parse_csv_list(value: str) -> List[str]:
        out: List[str] = []
        seen: Set[str] = set()
        for item in value.split(","):
            token = item.strip()
            if not token or token in seen:
                continue
            seen.add(token)
            out.append(token)
        return out

    @staticmethod
    def _parse_json_object(value: Any) -> Dict[str, Any]:
        if isinstance(value, dict):
            return value
        if isinstance(value, str):
            raw = value.strip()
            if not raw:
                return {}
            try:
                loaded = json.loads(raw)
                if isinstance(loaded, dict):
                    return loaded
            except Exception:
                return {}
        return {}

    @staticmethod
    def _sha256_json(value: Any) -> str:
        blob = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
        return hashlib.sha256(blob).hexdigest()

    def _resolve_dependency_signature(self, raw: Any) -> str:
        if isinstance(raw, str):
            return raw.strip()
        if raw is None:
            return ""
        return "sha256:" + self._sha256_json(self._safe_json(raw))

    @staticmethod
    def _context_order_signature(context_order: Any, args: Dict[str, Any]) -> str:
        if isinstance(context_order, str):
            return context_order.strip()
        if isinstance(context_order, (list, tuple)):
            return "->".join(str(item) for item in context_order)
        if isinstance(context_order, dict):
            return "->".join(str(key) for key in context_order.keys())
        if isinstance(args, dict):
            return "->".join(str(key) for key in args.keys())
        return ""

    def _resolve_observable_execution_state(self, merged: Dict[str, Any]) -> Dict[str, Any]:
        model_label = str(merged.get("model_label") or self.model_label or "").strip()
        provider_label = str(merged.get("provider_label") or self.provider_label or "").strip()
        prompt_template_version = str(
            merged.get("prompt_template_version") or self.prompt_template_version or ""
        ).strip()
        tool_schema_version = str(merged.get("tool_schema_version") or self.tool_schema_version or "").strip()

        decoding_raw = merged.get("decoding_params")
        if decoding_raw is None:
            decoding_raw = self.decoding_params
        decoding_params = self._parse_json_object(decoding_raw)

        decoding_params_signature = str(merged.get("decoding_params_signature") or "").strip()
        if not decoding_params_signature:
            decoding_params_signature = "sha256:" + self._sha256_json(
                {"decoding_params": self._safe_json(decoding_params)}
            )

        model_execution_signature = str(merged.get("model_execution_signature") or "").strip()
        if not model_execution_signature:
            model_execution_signature = "sha256:" + self._sha256_json(
                {
                    "model_label": model_label,
                    "provider_label": provider_label,
                    "prompt_template_version": prompt_template_version,
                    "tool_schema_version": tool_schema_version,
                    "decoding_params_signature": decoding_params_signature,
                }
            )
        behavior_probe_mode = str(
            merged.get("behavior_probe_mode") or self.behavior_probe_mode or "off"
        ).strip().lower()
        if behavior_probe_mode not in {"off", "observe", "enforce"}:
            behavior_probe_mode = "off"
        behavior_probe_suite_version = str(
            merged.get("behavior_probe_suite_version") or self.behavior_probe_suite_version or ""
        ).strip()
        behavior_probe_signature = str(
            merged.get("behavior_probe_signature") or self.behavior_probe_signature or ""
        ).strip()
        behavior_probe_checked_at = str(
            merged.get("behavior_probe_checked_at") or self.behavior_probe_checked_at or ""
        ).strip()

        return {
            "model_label": model_label,
            "provider_label": provider_label,
            "prompt_template_version": prompt_template_version,
            "tool_schema_version": tool_schema_version,
            "decoding_params": decoding_params,
            "decoding_params_signature": decoding_params_signature,
            "model_execution_signature": model_execution_signature,
            "behavior_probe_mode": behavior_probe_mode,
            "behavior_probe_suite_version": behavior_probe_suite_version,
            "behavior_probe_signature": behavior_probe_signature,
            "behavior_probe_checked_at": behavior_probe_checked_at,
        }

    def _resolve_authorization_state_vector_signature(
        self,
        merged: Dict[str, Any],
        *,
        policy_version: str,
        dependency_signature: str,
        context_ordering_signature: str,
        evidence_set_signature: str,
        evidence_order_signature: str,
        retrieval_params_signature: str,
        model_label: str,
        provider_label: str,
        prompt_template_version: str,
        tool_schema_version: str,
        decoding_params_signature: str,
        model_execution_signature: str,
        behavior_probe_mode: str,
        behavior_probe_suite_version: str,
        behavior_probe_signature: str,
    ) -> str:
        existing = str(merged.get("authorization_state_vector_signature") or "").strip()
        if existing:
            return existing
        return "sha256:" + self._sha256_json(
            {
                "policy_version": policy_version,
                "dependency_signature": dependency_signature,
                "context_ordering_signature": context_ordering_signature,
                "evidence_set_signature": evidence_set_signature,
                "evidence_order_signature": evidence_order_signature,
                "retrieval_params_signature": retrieval_params_signature,
                "model_label": model_label,
                "provider_label": provider_label,
                "prompt_template_version": prompt_template_version,
                "tool_schema_version": tool_schema_version,
                "decoding_params_signature": decoding_params_signature,
                "model_execution_signature": model_execution_signature,
                "behavior_probe_mode": behavior_probe_mode,
                "behavior_probe_suite_version": behavior_probe_suite_version,
                "behavior_probe_signature": behavior_probe_signature,
            }
        )

    def _resolve_evidence_signatures(
        self,
        merged: Dict[str, Any],
        policy_version: str,
        dependency_signature: str,
        semantic_fields: Any,
        context_ordering_signature: str,
    ) -> Dict[str, str]:
        evidence_set_signature = str(merged.get("evidence_set_signature") or "").strip()
        if not evidence_set_signature:
            evidence_set_signature = "sha256:" + self._sha256_json(
                {"semantic_fields": self._safe_json(semantic_fields)}
            )

        evidence_order_signature = str(merged.get("evidence_order_signature") or "").strip()
        if not evidence_order_signature:
            order_material = merged.get("context_order")
            if order_material is None:
                order_material = context_ordering_signature
            evidence_order_signature = "sha256:" + self._sha256_json(
                {"context_order": self._safe_json(order_material)}
            )

        retrieval_params_signature = str(merged.get("retrieval_params_signature") or "").strip()
        if not retrieval_params_signature:
            retrieval_material = merged.get("retrieval_params")
            if retrieval_material is None:
                retrieval_material = {"context_mode": self.context_mode}
            retrieval_params_signature = "sha256:" + self._sha256_json(
                {"retrieval_params": self._safe_json(retrieval_material)}
            )

        authorization_state_vector_signature = str(
            merged.get("authorization_state_vector_signature") or ""
        ).strip()
        if not authorization_state_vector_signature:
            authorization_state_vector_signature = "sha256:" + self._sha256_json(
                {
                    "policy_version": policy_version,
                    "dependency_signature": dependency_signature,
                    "context_ordering_signature": context_ordering_signature,
                    "evidence_set_signature": evidence_set_signature,
                    "evidence_order_signature": evidence_order_signature,
                    "retrieval_params_signature": retrieval_params_signature,
                }
            )

        return {
            "evidence_set_signature": evidence_set_signature,
            "evidence_order_signature": evidence_order_signature,
            "retrieval_params_signature": retrieval_params_signature,
            "authorization_state_vector_signature": authorization_state_vector_signature,
        }

    @classmethod
    def _is_noisy_session_key(cls, key: str) -> bool:
        raw = key.strip().lower()
        if raw in cls._DEFAULT_NOISY_SESSION_KEYS:
            return True
        return raw.endswith("_ts") or raw.endswith("_time")

    def _sanitize_session_state(self, value: Any, depth: int = 0) -> Any:
        if depth >= 5:
            return "<truncated>"
        if isinstance(value, dict):
            out: Dict[str, Any] = {}
            keys = sorted(str(k) for k in value.keys())
            for key in keys[:64]:
                if self._is_noisy_session_key(key):
                    continue
                out[key] = self._sanitize_session_state(value.get(key), depth + 1)
            return out
        if isinstance(value, list):
            return [self._sanitize_session_state(item, depth + 1) for item in value[:32]]
        if isinstance(value, tuple):
            return [self._sanitize_session_state(item, depth + 1) for item in value[:32]]
        if isinstance(value, (str, int, float, bool)) or value is None:
            return value
        return str(value)

    def _extract_adk_session_state(self, tool_context: Any) -> Dict[str, Any]:
        inv = getattr(tool_context, "_invocation_context", None) or getattr(tool_context, "invocation_context", None)
        if inv is None:
            return {}
        session = getattr(inv, "session", None)
        if session is None:
            return {}
        state = getattr(session, "state", None)
        if not isinstance(state, dict):
            return {}
        return self._safe_json(state)

    def _extract_session_signature(self, session_state: Dict[str, Any]) -> str:
        if self.session_signature_mode == "off":
            return ""
        if not isinstance(session_state, dict) or not session_state:
            return ""
        material: Any
        if self.session_signature_mode == "full":
            material = session_state
        else:
            material = self._sanitize_session_state(session_state)
        if not material:
            return ""
        return "sha256:" + self._sha256_json(material)

    @staticmethod
    def _default_identity_fields(args: Dict[str, Any]) -> Dict[str, Any]:
        out: Dict[str, Any] = {}
        if not isinstance(args, dict):
            return out
        for key, value in args.items():
            if key == "id" or key.endswith("_id") or key in {"order_id", "customer_id", "payment_id", "policy_id"}:
                out[str(key)] = value
        return out

    @staticmethod
    def _path_tokens(path: str) -> List[Any]:
        tokens: List[Any] = []
        buf = ""
        i = 0
        while i < len(path):
            ch = path[i]
            if ch == ".":
                if buf:
                    tokens.append(buf)
                    buf = ""
                i += 1
                continue
            if ch == "[":
                if buf:
                    tokens.append(buf)
                    buf = ""
                end = path.find("]", i + 1)
                if end < 0:
                    break
                raw = path[i + 1 : end].strip().strip('"').strip("'")
                if raw.isdigit():
                    tokens.append(int(raw))
                elif raw:
                    tokens.append(raw)
                i = end + 1
                continue
            buf += ch
            i += 1
        if buf:
            tokens.append(buf)
        return [token for token in tokens if token != ""]

    @classmethod
    def _extract_value_by_path(cls, root: Any, path: str) -> Tuple[Any, bool]:
        current = root
        tokens = cls._path_tokens(path)
        if not tokens:
            return None, False
        for token in tokens:
            if isinstance(token, int):
                if not isinstance(current, (list, tuple)) or token < 0 or token >= len(current):
                    return None, False
                current = current[token]
                continue
            if isinstance(current, dict) and token in current:
                current = current[token]
                continue
            return None, False
        return current, True

    def _extract_semantic_fields_from_paths(self, source: Dict[str, Any], paths: List[str]) -> Dict[str, Any]:
        out: Dict[str, Any] = {}
        for raw_path in paths:
            path = str(raw_path or "").strip()
            if not path:
                continue
            value, found = self._extract_value_by_path(source, path)
            if found:
                out[path] = self._safe_json(value)
            else:
                out[path] = {"_kyb_missing_path": path}
        return out

    def _build_default_intent_context(self, tool_name: str, args: Dict[str, Any], tool_context: Any) -> Dict[str, Any]:
        safe_args = self._safe_json(args if isinstance(args, dict) else {"value": args})
        session_state = self._extract_adk_session_state(tool_context)
        session_signature = self._extract_session_signature(session_state)

        inv = getattr(tool_context, "_invocation_context", None) or getattr(tool_context, "invocation_context", None)
        invocation_id = str(getattr(inv, "invocation_id", "") or "") if inv is not None else ""
        session_obj = getattr(inv, "session", None) if inv is not None else None
        session_id = str(getattr(session_obj, "id", "") or "") if session_obj is not None else ""
        function_call_id = str(getattr(tool_context, "function_call_id", "") or "")

        source = {
            "tool_name": tool_name,
            "tool_args": safe_args,
            "session_state": session_state,
            "session_signature": session_signature,
            "invocation": {
                "invocation_id": invocation_id,
                "session_id": session_id,
                "function_call_id": function_call_id,
            },
        }
        selected_paths = list(self.semantic_paths)
        semantic_fields = self._extract_semantic_fields_from_paths(source, selected_paths)
        if "tool_name" not in semantic_fields:
            semantic_fields["tool_name"] = tool_name
        if session_signature and "session_signature" not in semantic_fields:
            semantic_fields["session_signature"] = session_signature

        policy_version = self.policy_version
        dependency_raw: Any = self.dependency_signature
        if not dependency_raw:
            dependency_raw = {
                "session_signature": session_signature,
                "semantic_paths": selected_paths,
            }
        dependency_sig = self._resolve_dependency_signature(dependency_raw)
        context_order_sig = "->".join(selected_paths)
        identity_fields = self._default_identity_fields(safe_args if isinstance(safe_args, dict) else {})
        merged_for_signatures = {
            "context_order": selected_paths,
            "retrieval_params": {
                "source": "adk_default_context",
                "semantic_paths": selected_paths,
            },
        }
        signatures = self._resolve_evidence_signatures(
            merged=merged_for_signatures,
            policy_version=policy_version,
            dependency_signature=dependency_sig,
            semantic_fields=semantic_fields,
            context_ordering_signature=context_order_sig,
        )
        execution_state = self._resolve_observable_execution_state(merged_for_signatures)
        signatures["authorization_state_vector_signature"] = self._resolve_authorization_state_vector_signature(
            merged=merged_for_signatures,
            policy_version=policy_version,
            dependency_signature=dependency_sig,
            context_ordering_signature=context_order_sig,
            evidence_set_signature=signatures.get("evidence_set_signature", ""),
            evidence_order_signature=signatures.get("evidence_order_signature", ""),
            retrieval_params_signature=signatures.get("retrieval_params_signature", ""),
            model_label=execution_state.get("model_label", ""),
            provider_label=execution_state.get("provider_label", ""),
            prompt_template_version=execution_state.get("prompt_template_version", ""),
            tool_schema_version=execution_state.get("tool_schema_version", ""),
            decoding_params_signature=execution_state.get("decoding_params_signature", ""),
            model_execution_signature=execution_state.get("model_execution_signature", ""),
            behavior_probe_mode=execution_state.get("behavior_probe_mode", "off"),
            behavior_probe_suite_version=execution_state.get("behavior_probe_suite_version", ""),
            behavior_probe_signature=execution_state.get("behavior_probe_signature", ""),
        )
        fingerprint_material = {
            "tool": tool_name,
            "policy_version": policy_version,
            "dependency_signature": dependency_sig,
            "semantic_fields": semantic_fields,
            "context_ordering_signature": context_order_sig,
            **signatures,
            **execution_state,
        }
        fingerprint_hash = "sha256:" + self._sha256_json(fingerprint_material)
        return {
            "policy_version": policy_version,
            "dependency_signature": dependency_sig,
            "semantic_fields": semantic_fields,
            "context_order": selected_paths,
            "context_ordering_signature": context_order_sig,
            "identity_fields": identity_fields,
            "fingerprint_hash": fingerprint_hash,
            **signatures,
            **execution_state,
        }

    def _build_tool_intent(self, tool_name: str, args: Dict[str, Any], tool_context: Any) -> Dict[str, Any]:
        base_context: Dict[str, Any] = {}
        if self.context_mode == "adk_default":
            base_context = self._build_default_intent_context(tool_name, args, tool_context)

        provided: Dict[str, Any] = {}
        if self.intent_context_provider is not None and self.context_mode == "provider":
            try:
                candidate = self.intent_context_provider(tool_name, args, tool_context)
                if isinstance(candidate, dict):
                    provided = candidate
            except Exception as exc:
                provided = {"provider_error": str(exc)}

        merged: Dict[str, Any] = dict(base_context)
        for key, value in provided.items():
            if value is not None:
                merged[key] = value

        policy_version = str(merged.get("policy_version") or self.policy_version or "")
        dep_raw = merged.get("dependency_signature")
        if dep_raw is None:
            dep_raw = self.dependency_signature
        dependency_sig = self._resolve_dependency_signature(dep_raw)
        semantic_fields = merged.get("semantic_fields")
        context_order = merged.get("context_order")
        context_order_sig = str(merged.get("context_ordering_signature") or "").strip()
        if not context_order_sig:
            context_order_sig = self._context_order_signature(context_order, args)
        identity_fields = merged.get("identity_fields")
        intent_key = str(merged.get("intent_key") or "")
        if not intent_key:
            intent_key, identity_fields = build_intent_key(tool_name, args, identity_fields)
        field_hashes = merged.get("field_semantic_hashes")
        if not isinstance(field_hashes, dict):
            field_hashes = merged.get("field_hashes")
        if not isinstance(field_hashes, dict):
            field_hashes = {}
        signatures = self._resolve_evidence_signatures(
            merged=merged,
            policy_version=policy_version,
            dependency_signature=dependency_sig,
            semantic_fields=semantic_fields if semantic_fields is not None else args,
            context_ordering_signature=context_order_sig,
        )
        execution_state = self._resolve_observable_execution_state(merged)
        signatures["authorization_state_vector_signature"] = self._resolve_authorization_state_vector_signature(
            merged=merged,
            policy_version=policy_version,
            dependency_signature=dependency_sig,
            context_ordering_signature=context_order_sig,
            evidence_set_signature=signatures.get("evidence_set_signature", ""),
            evidence_order_signature=signatures.get("evidence_order_signature", ""),
            retrieval_params_signature=signatures.get("retrieval_params_signature", ""),
            model_label=execution_state.get("model_label", ""),
            provider_label=execution_state.get("provider_label", ""),
            prompt_template_version=execution_state.get("prompt_template_version", ""),
            tool_schema_version=execution_state.get("tool_schema_version", ""),
            decoding_params_signature=execution_state.get("decoding_params_signature", ""),
            model_execution_signature=execution_state.get("model_execution_signature", ""),
            behavior_probe_mode=execution_state.get("behavior_probe_mode", "off"),
            behavior_probe_suite_version=execution_state.get("behavior_probe_suite_version", ""),
            behavior_probe_signature=execution_state.get("behavior_probe_signature", ""),
        )
        approval_gate_id = str(merged.get("approval_gate_id") or "").strip()
        intent_out = {
            "intent_key": intent_key,
            "policy_version": policy_version,
            "dependency_signature": dependency_sig,
            "semantic_fields": semantic_fields if semantic_fields is not None else args,
            "context_ordering_signature": context_order_sig,
            "field_hashes": field_hashes,
            "fingerprint_hash": str(merged.get("fingerprint_hash") or ""),
            "identity_fields": identity_fields if isinstance(identity_fields, dict) else {},
            **signatures,
            **execution_state,
        }
        if approval_gate_id:
            intent_out["approval_gate_id"] = approval_gate_id
        return intent_out

    def _wrap_adk_toolunion(self, tool: Any) -> Any:
        # Callable (function tool)
        if callable(tool) and not _is_adk_basetool(tool) and not self._is_adk_toolset(tool):
            return self._wrap_callable_tool(tool)

        # BaseTool instance
        if _is_adk_basetool(tool):
            return self._wrap_basetool(tool)

        # BaseToolset instance
        if self._is_adk_toolset(tool):
            return self._wrap_toolset(tool)

        return tool
    
    def _is_adk_toolset(self, obj: Any) -> bool:
        # duck-typing: toolset has get_tools_with_prefix/get_tools
        return hasattr(obj, "get_tools") and hasattr(obj, "get_tools_with_prefix")

    

    def _wrap_callable_tool(self, fn: Any) -> Any:
        if _mark(fn, "_kyb_wrapped"):
            return fn

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            run_id = kyb_run_id.get()
            if not run_id:
                # No tool_context for plain callables; do not guess. Either:
                # 1) raise, or 2) just call fn without emitting.
                raise RuntimeError("kyb_run_id not set; callable tool invoked outside ADK tool_context")
            tool_name = getattr(fn, "__name__", "tool")
            print(f"[kybernis] tool.call run_id={run_id} tool={tool_name}", flush=True)
            client = kyb_client()

            # pre: emit
            fire_and_forget(
                lambda client=client, run_id=run_id, tool_name=tool_name: client.emit_event(
                    run_id, "tool.call", {"tool": tool_name}
                )
            )

            try:
                out = fn(*args, **kwargs)
                fire_and_forget(
                    lambda client=client, run_id=run_id, tool_name=tool_name: client.emit_event(
                        run_id, "tool.result", {"tool": tool_name, "ok": True}
                    )
                )
                return out
            except Exception as exc:
                err_text = str(exc)
                fire_and_forget(
                    lambda client=client, run_id=run_id, tool_name=tool_name, err_text=err_text: client.emit_event(
                        run_id, "tool.error", {"tool": tool_name, "error": err_text}
                    )
                )
                raise

        _mark(wrapper, "_kyb_wrapped")
        return wrapper

    def _wrap_basetool(self, tool: Any) -> Any:
        if _mark(tool, "_kyb_run_async_wrapped"):
            return tool

        original = tool.run_async  # bound

        @functools.wraps(original)
        async def run_async_wrapped(*, args: dict, tool_context: Any) -> Any:
            # Ensure run context exists (it should from agent patch, but tools can be invoked independently)
            inv_ctx = getattr(tool_context, "_invocation_context", None) or getattr(tool_context, "invocation_context", None)

            inv_id = ""
            if inv_ctx is None:
                print("[kybernis] WARNING: inv_ctx is None", flush=True)
            else:
                inv_id = str(getattr(inv_ctx, "invocation_id", "") or "").strip()

            # Prefer existing bound run_id; only resolve if missing
            run_id = kyb_run_id.get() or ""
            if not run_id and inv_id:
                run_id = resolve_run_id_from_invocation_id(inv_id)

            if not run_id:
                raise RuntimeError("missing invocation_id and kyb_run_id; cannot emit")

            token_run = kyb_run_id.set(run_id)

            client = kyb_client()
            tool_name = getattr(tool, "name", tool.__class__.__name__)
            call_id = str(getattr(tool_context, "function_call_id", "") or "")
            safe_args = self._safe_json(args if isinstance(args, dict) else {"value": args})
            intent = self._build_tool_intent(tool_name, safe_args if isinstance(safe_args, dict) else {"value": safe_args}, tool_context)
            intent_key = str(intent.get("intent_key") or f"intent:{uuid.uuid4().hex}")
            fingerprint_hash = str(intent.get("fingerprint_hash") or "")
            task_run_id = intent_key
            attrs = {
                "tool": {
                    "tool_name": tool_name,
                    "function_call_id": call_id,
                    "intent_key": intent_key,
                    "policy_version": intent.get("policy_version", ""),
                    "dependency_signature": intent.get("dependency_signature", ""),
                    "context_ordering_signature": intent.get("context_ordering_signature", ""),
                    "evidence_set_signature": intent.get("evidence_set_signature", ""),
                    "evidence_order_signature": intent.get("evidence_order_signature", ""),
                    "retrieval_params_signature": intent.get("retrieval_params_signature", ""),
                    "authorization_state_vector_signature": intent.get("authorization_state_vector_signature", ""),
                    "model_label": intent.get("model_label", ""),
                    "provider_label": intent.get("provider_label", ""),
                    "prompt_template_version": intent.get("prompt_template_version", ""),
                    "tool_schema_version": intent.get("tool_schema_version", ""),
                    "decoding_params": intent.get("decoding_params", {}),
                    "decoding_params_signature": intent.get("decoding_params_signature", ""),
                    "model_execution_signature": intent.get("model_execution_signature", ""),
                    "behavior_probe_mode": intent.get("behavior_probe_mode", "off"),
                    "behavior_probe_suite_version": intent.get("behavior_probe_suite_version", ""),
                    "behavior_probe_signature": intent.get("behavior_probe_signature", ""),
                    "behavior_probe_checked_at": intent.get("behavior_probe_checked_at", ""),
                }
            }
            intent_payload = {
                "tool_name": tool_name,
                "policy_version": intent.get("policy_version", ""),
                "dependency_signature": intent.get("dependency_signature", ""),
                "context_ordering_signature": intent.get("context_ordering_signature", ""),
                "evidence_set_signature": intent.get("evidence_set_signature", ""),
                "evidence_order_signature": intent.get("evidence_order_signature", ""),
                "retrieval_params_signature": intent.get("retrieval_params_signature", ""),
                "authorization_state_vector_signature": intent.get("authorization_state_vector_signature", ""),
                "model_label": intent.get("model_label", ""),
                "provider_label": intent.get("provider_label", ""),
                "prompt_template_version": intent.get("prompt_template_version", ""),
                "tool_schema_version": intent.get("tool_schema_version", ""),
                "decoding_params": intent.get("decoding_params", {}),
                "decoding_params_signature": intent.get("decoding_params_signature", ""),
                "model_execution_signature": intent.get("model_execution_signature", ""),
                "behavior_probe_mode": intent.get("behavior_probe_mode", "off"),
                "behavior_probe_suite_version": intent.get("behavior_probe_suite_version", ""),
                "behavior_probe_signature": intent.get("behavior_probe_signature", ""),
                "behavior_probe_checked_at": intent.get("behavior_probe_checked_at", ""),
                "identity_fields": intent.get("identity_fields", {}),
                "semantic_fields": intent.get("semantic_fields", {}),
            }
            if intent.get("approval_gate_id"):
                intent_payload["approval_gate_id"] = intent.get("approval_gate_id")
            if isinstance(intent.get("field_hashes"), dict) and intent.get("field_hashes"):
                intent_payload["field_semantic_hashes"] = intent.get("field_hashes", {})
            mutation_id = ""
            runtime_intent_hash = ""
            guard_blocked = False
            task_run_started = False

            try:
                # Start the task-run row first so artifact/finish calls never race it.
                print(f"[kybernis] tool.enter name={tool_name} run_id={run_id}", flush=True)
                try:
                    client.task_run_start(
                        run_id=run_id,
                        task_run_id=task_run_id,
                        kind="tool",
                        name=tool_name,
                        idempotency_key=intent_key,
                        attributes=attrs,
                    )
                    task_run_started = True
                except Exception as start_exc:
                    print(f"[kybernis] task_run_start failed: {start_exc}", flush=True)
                if task_run_started:
                    fire_and_forget(
                        lambda client=client, run_id=run_id, task_run_id=task_run_id, intent_key=intent_key, fingerprint_hash=fingerprint_hash, intent=intent: client.task_run_artifact(
                            run_id=run_id,
                            task_run_id=task_run_id,
                            kind="intent_context",
                            payload={
                                "intent_key": intent_key,
                                "fingerprint_hash": fingerprint_hash,
                                "policy_version": intent.get("policy_version", ""),
                                "dependency_signature": intent.get("dependency_signature", ""),
                                "context_ordering_signature": intent.get("context_ordering_signature", ""),
                                "evidence_set_signature": intent.get("evidence_set_signature", ""),
                                "evidence_order_signature": intent.get("evidence_order_signature", ""),
                                "retrieval_params_signature": intent.get("retrieval_params_signature", ""),
                                "authorization_state_vector_signature": intent.get("authorization_state_vector_signature", ""),
                                "model_label": intent.get("model_label", ""),
                                "provider_label": intent.get("provider_label", ""),
                                "prompt_template_version": intent.get("prompt_template_version", ""),
                                "tool_schema_version": intent.get("tool_schema_version", ""),
                                "decoding_params_signature": intent.get("decoding_params_signature", ""),
                                "model_execution_signature": intent.get("model_execution_signature", ""),
                                "decoding_params": intent.get("decoding_params", {}),
                                "behavior_probe_mode": intent.get("behavior_probe_mode", "off"),
                                "behavior_probe_suite_version": intent.get("behavior_probe_suite_version", ""),
                                "behavior_probe_signature": intent.get("behavior_probe_signature", ""),
                                "behavior_probe_checked_at": intent.get("behavior_probe_checked_at", ""),
                                "semantic_fields": intent.get("semantic_fields", {}),
                                "identity_fields": intent.get("identity_fields", {}),
                            },
                        )
                    )
                fire_and_forget(lambda: client.emit_event(run_id, "tool.call", {"tool": tool_name, "args": safe_args, "intent_key": intent_key}))

                if self.drift_guard_enabled:
                    approval_attempts = 0
                    while True:
                        guard = client.intent_guard(
                            run_id=run_id,
                            idempotency_key=intent_key,
                            intent=intent_payload,
                            original_payload={"intent": intent_payload, "semantic_fields": intent.get("semantic_fields", {}), "args": safe_args},
                            requested_mode=self.drift_guard_mode,
                        )
                        mutation_id = str(guard.get("mutation_id") or "")
                        runtime_intent_hash = str(guard.get("runtime_intent_hash") or "")
                        if bool(guard.get("drift_detected")):
                            current_field_hashes = guard.get("field_semantic_hashes", {})
                            if not isinstance(current_field_hashes, dict):
                                current_field_hashes = {}
                            drift_payload = {
                                "tool": tool_name,
                                "intent_key": intent_key,
                                "drift_type": "type_b_context_mutation",
                                "drift_fields_changed": guard.get("drift_fields_changed", []),
                                "decision": guard.get("decision", ""),
                                "policy_mode": guard.get("policy_mode", ""),
                                "current_field_hashes": current_field_hashes,
                                "policy_version": intent.get("policy_version", ""),
                                "dependency_signature": intent.get("dependency_signature", ""),
                                "context_ordering_signature": intent.get("context_ordering_signature", ""),
                                "evidence_set_signature": intent.get("evidence_set_signature", ""),
                                "evidence_order_signature": intent.get("evidence_order_signature", ""),
                                "retrieval_params_signature": intent.get("retrieval_params_signature", ""),
                                "authorization_state_vector_signature": intent.get("authorization_state_vector_signature", ""),
                                "model_label": intent.get("model_label", ""),
                                "provider_label": intent.get("provider_label", ""),
                                "prompt_template_version": intent.get("prompt_template_version", ""),
                                "tool_schema_version": intent.get("tool_schema_version", ""),
                                "decoding_params": intent.get("decoding_params", {}),
                                "decoding_params_signature": intent.get("decoding_params_signature", ""),
                                "model_execution_signature": intent.get("model_execution_signature", ""),
                                "behavior_probe_mode": intent.get("behavior_probe_mode", "off"),
                                "behavior_probe_suite_version": intent.get("behavior_probe_suite_version", ""),
                                "behavior_probe_signature": intent.get("behavior_probe_signature", ""),
                                "behavior_probe_checked_at": intent.get("behavior_probe_checked_at", ""),
                            }
                            fire_and_forget(lambda: client.emit_event(run_id, "DRIFT_DETECTED", drift_payload))
                            if task_run_started:
                                fire_and_forget(
                                    lambda client=client, run_id=run_id, task_run_id=task_run_id, drift_payload=drift_payload: client.task_run_artifact(
                                        run_id=run_id,
                                        task_run_id=task_run_id,
                                        kind="drift_detected",
                                        payload=drift_payload,
                                    )
                                )
                        if bool(guard.get("block_execute")):
                            reason = str(guard.get("reason") or "drift_guard_blocked")
                            gate_id = str(guard.get("gate_id") or "")
                            gate_status = ""
                            if reason == "approval_required" and gate_id:
                                approval_attempts += 1
                                gate_timeout_s = int(float(os.getenv("KYB_GATE_TIMEOUT_S", "300") or "300"))
                                gate_poll_s = float(os.getenv("KYB_GATE_POLL_S", "1") or "1")
                                gate_state = client.gate_wait(
                                    gate_id=gate_id,
                                    timeout_s=gate_timeout_s,
                                    poll_interval_s=gate_poll_s,
                                    run_id=run_id,
                                )
                                gate_data = (gate_state.get("data", {}) or {}) if isinstance(gate_state, dict) else {}
                                gate_status = str(gate_data.get("status") or "").upper()
                                if gate_status in {"APPROVED", "CONSUMED", "OVERRIDDEN"}:
                                    intent_payload["approval_gate_id"] = gate_id
                                    intent["approval_gate_id"] = gate_id
                                    continue
                                if gate_status == "DENIED":
                                    reason = "approval_denied"
                                elif approval_attempts >= 3:
                                    reason = "approval_retry_exhausted"
                                else:
                                    reason = "approval_timeout"
                            guard_blocked = True
                            block_decision = str(guard.get("decision") or "")
                            block_mode = str(guard.get("policy_mode") or self.drift_guard_mode or "")
                            blocked_result = {
                                "status": "error",
                                "message": reason,
                                "error_code": "DRIFT_GUARD_BLOCKED",
                                "kybernis": {
                                    "blocked": True,
                                    "reason": reason,
                                    "decision": block_decision,
                                    "mode": block_mode,
                                    "intent_key": intent_key,
                                    "mutation_id": mutation_id,
                                    "drift_fields_changed": guard.get("drift_fields_changed", []),
                                    "gate_id": gate_id,
                                    "gate_status": gate_status,
                                },
                            }
                            fire_and_forget(
                                lambda client=client, run_id=run_id, tool_name=tool_name, reason=reason, intent_key=intent_key, block_decision=block_decision, block_mode=block_mode: client.emit_event(
                                    run_id,
                                    "tool.blocked",
                                    {
                                        "tool": tool_name,
                                        "reason": reason,
                                        "intent_key": intent_key,
                                        "decision": block_decision,
                                        "mode": block_mode,
                                    },
                                )
                            )
                            if task_run_started:
                                fire_and_forget(
                                    lambda client=client, run_id=run_id, task_run_id=task_run_id, blocked_result=blocked_result: client.task_run_artifact(
                                        run_id=run_id,
                                        task_run_id=task_run_id,
                                        kind="tool_blocked",
                                        payload=blocked_result,
                                    )
                                )
                            if task_run_started:
                                fire_and_forget(
                                    lambda client=client, run_id=run_id, task_run_id=task_run_id, reason=reason: client.task_run_finish(
                                        run_id=run_id,
                                        task_run_id=task_run_id,
                                        status="BLOCKED",
                                        error=reason,
                                        error_code="DRIFT_GUARD",
                                    )
                                )
                            return blocked_result
                        break

                with kyb_span(
                    client=client,
                    span_type="tool",
                    name=tool_name,
                    attributes=attrs,
                    idempotency_key=intent_key,
                    input_hash=runtime_intent_hash or fingerprint_hash or None,
                ):
                    res = await original(args=args, tool_context=tool_context)
                fire_and_forget(lambda: client.emit_event(run_id, "tool.result", {"tool": tool_name, "ok": True}))
                if task_run_started:
                    fire_and_forget(
                        lambda client=client, run_id=run_id, task_run_id=task_run_id, res=res: client.task_run_artifact(
                            run_id=run_id,
                            task_run_id=task_run_id,
                            kind="tool_result",
                            payload={"ok": True, "result": self._safe_json(res)},
                        )
                    )
                    fire_and_forget(
                        lambda client=client, run_id=run_id, task_run_id=task_run_id: client.task_run_finish(
                            run_id=run_id,
                            task_run_id=task_run_id,
                            status="OK",
                            result={"ok": True},
                        )
                    )
                if mutation_id:
                    fire_and_forget(
                        lambda: client.intent_commit(
                            mutation_id=mutation_id,
                            runtime_intent_hash=runtime_intent_hash,
                            receipt={"tool": tool_name, "status": "ok"},
                            provider="adk_tool",
                        )
                    )
                return res
            except Exception as exc:
                err_text = str(exc)
                err_code = type(exc).__name__
                fire_and_forget(
                    lambda client=client, run_id=run_id, tool_name=tool_name, err_text=err_text: client.emit_event(
                        run_id, "tool.error", {"tool": tool_name, "error": err_text}
                    )
                )
                if task_run_started:
                    fire_and_forget(
                        lambda client=client, run_id=run_id, task_run_id=task_run_id, err_text=err_text: client.task_run_artifact(
                            run_id=run_id,
                            task_run_id=task_run_id,
                            kind="tool_error",
                            payload={"error": err_text},
                        )
                    )
                    fire_and_forget(
                        lambda client=client, run_id=run_id, task_run_id=task_run_id, err_text=err_text, err_code=err_code: client.task_run_finish(
                            run_id=run_id,
                            task_run_id=task_run_id,
                            status="ERROR",
                            error=err_text,
                            error_code=err_code,
                        )
                    )
                if mutation_id and not guard_blocked:
                    if self._is_ambiguous_execution_error(exc):
                        unknown_reason = self._unknown_outcome_reason(exc)
                        fire_and_forget(
                            lambda client=client, run_id=run_id, tool_name=tool_name, mutation_id=mutation_id, unknown_reason=unknown_reason: client.emit_event(
                                run_id,
                                "tool.outcome_unknown",
                                {"tool": tool_name, "mutation_id": mutation_id, "reason": unknown_reason},
                            )
                        )
                        if task_run_started:
                            fire_and_forget(
                                lambda client=client, run_id=run_id, task_run_id=task_run_id, mutation_id=mutation_id, unknown_reason=unknown_reason: client.task_run_artifact(
                                    run_id=run_id,
                                    task_run_id=task_run_id,
                                    kind="tool_outcome_unknown",
                                    payload={"mutation_id": mutation_id, "reason": unknown_reason},
                                )
                            )
                        fire_and_forget(
                            lambda client=client, mutation_id=mutation_id, unknown_reason=unknown_reason: client.intent_unknown(
                                mutation_id=mutation_id,
                                reason=unknown_reason,
                            )
                        )
                    else:
                        fire_and_forget(
                            lambda client=client, mutation_id=mutation_id, err_text=err_text: client.intent_fail(
                                mutation_id=mutation_id, reason=err_text
                            )
                        )
                raise
            finally:
                kyb_run_id.reset(token_run)

        try:
            tool.run_async = run_async_wrapped
        except Exception:
            object.__setattr__(tool, "run_async", run_async_wrapped)

        return tool

    def _wrap_toolset(self, toolset: Any) -> Any:
        if _mark(toolset, "_kyb_toolset_wrapped"):
            return toolset

        # wrap get_tools
        if hasattr(toolset, "get_tools"):
            original_get_tools = toolset.get_tools

            @functools.wraps(original_get_tools)
            async def get_tools_wrapped(readonly_context: Any = None):
                tools = await original_get_tools(readonly_context)
                return [self._wrap_basetool(t) for t in tools]

            try:
                toolset.get_tools = get_tools_wrapped
            except Exception:
                object.__setattr__(toolset, "get_tools", get_tools_wrapped)

        # wrap get_tools_with_prefix (if toolset uses it heavily)
        if hasattr(toolset, "get_tools_with_prefix"):
            original_get_tools_with_prefix = toolset.get_tools_with_prefix

            @functools.wraps(original_get_tools_with_prefix)
            async def get_tools_with_prefix_wrapped(readonly_context: Any = None):
                tools = await original_get_tools_with_prefix(readonly_context)
                return [self._wrap_basetool(t) for t in tools]

            try:
                toolset.get_tools_with_prefix = get_tools_with_prefix_wrapped
            except Exception:
                object.__setattr__(toolset, "get_tools_with_prefix", get_tools_with_prefix_wrapped)

        return toolset

    def _wrap_adk_tools(self, agent: Any) -> None:
        print("Wrapping ADK tools...", agent)
        tools = getattr(agent, "tools", None)
        if not isinstance(tools, list):
            return

        updated = False
        wrapped_tools = []

        for t in tools:
            wt = self._wrap_adk_toolunion(t)
            wrapped_tools.append(wt)
            if wt is not t:
                updated = True

            nested_agent = getattr(t, "agent", None)
            if nested_agent is not None:
                self._wrap_adk_tools(nested_agent)

        if updated:
            try:
                setattr(agent, "tools", wrapped_tools)
            except Exception:
                object.__setattr__(agent, "tools", wrapped_tools)



    def _intercept_many(self, obj: Any, method_names: list, framework: str):
        for name in method_names:
            if hasattr(obj, name):
                self._intercept(obj, name, framework=framework)
        if self.auto_intercept:
            self._intercept_all_methods(obj, framework=framework)
        return obj

    def _intercept_all_methods(self, obj: Any, framework: str):
        for name in dir(obj):
            if name.startswith("_"):
                continue
            if self.allow_methods and name not in self.allow_methods:
                continue
            if name in self.exclude_methods:
                continue
            try:
                attr = getattr(obj, name)
            except Exception:
                continue
            if not (inspect.ismethod(attr) or inspect.isfunction(attr) or inspect.iscoroutinefunction(attr)):
                continue
            self._intercept(obj, name, framework=framework)
        return obj

    def _intercept(self, obj: Any, method_name: str, framework: str):
        original_method = getattr(obj, method_name)
        if getattr(original_method, "_kyb_wrapped", False):
            return obj

        if inspect.iscoroutinefunction(original_method):
            @functools.wraps(original_method)
            async def wrapper(*args, **kwargs):
                task_id = kwargs.get("task_id") or f"task_{uuid.uuid4().hex}"
                run_id = self._ensure_run_id(framework, obj, args, kwargs)
                await self._run_async(self._check_protocol_compliance, method_name, framework, task_id, run_id)
                await self._run_async(self._maybe_invoke_mcp, method_name, framework, obj, args, kwargs, run_id)
                if framework == "google_adk":
                    print("Calling original method in async wrapper for ADK tool")
                    result = await original_method(*args, **kwargs)
                else:
                    result = await original_method(*args, **kwargs)
                await self._run_async(self._report_task_completion, method_name, framework, task_id, run_id, result, kwargs)
                return result
        else:
            @functools.wraps(original_method)
            def wrapper(*args, **kwargs):
                task_id = kwargs.get("task_id") or f"task_{uuid.uuid4().hex}"
                run_id = self._ensure_run_id(framework, obj, args, kwargs)
                self._check_protocol_compliance(method_name, framework, task_id, run_id)
                self._maybe_invoke_mcp(method_name, framework, obj, args, kwargs, run_id)
                if framework == "google_adk":
                    print("Calling original method in sync wrapper for ADK tool")
                    result = original_method(*args, **kwargs)
                else:
                    result = original_method(*args, **kwargs)
                self._report_task_completion(method_name, framework, task_id, run_id, result, kwargs)
                return result

        setattr(wrapper, "_kyb_wrapped", True)
        setattr(obj, method_name, wrapper)
        return obj

    async def _run_async(self, fn: Callable[..., Any], *args, **kwargs) -> Any:
        try:
            to_thread = asyncio.to_thread
        except AttributeError:
            loop = asyncio.get_running_loop()
            return await loop.run_in_executor(None, functools.partial(fn, *args, **kwargs))
        return await to_thread(fn, *args, **kwargs)

    def _resolve_mode(self, mode: str) -> Tuple[bool, bool]:
        raw = (mode or os.getenv("KYB_SDK_MODE", "")).strip().lower()
        if raw in {"proxy", "proxy-first", "a2a+proxy"}:
            return True, False
        if raw in {"mcp", "mcp-first", "a2a+mcp"}:
            return False, True
        if raw in {"hybrid", "both", "a2a+proxy+mcp"}:
            return True, True
        return False, False

    def _check_protocol_compliance(self, method_name: str, framework: str, task_id: str, run_id: str) -> None:
        if not self.gateway_url:
            return
        payload = {
            "jsonrpc": "2.0",
            "id": str(uuid.uuid4()),
            "method": "tasks/stream",
            "params": {
                "task_id": task_id,
                "framework": framework,
                "method": method_name,
                "run_id": run_id,
                "tenant_id": self.tenant_id,
            },
        }
        resp = self._post_a2a(payload)
        if resp is None:
            return
        if resp.status_code == 429:
            raise KybernisProtocolException("rate_limited")
        if resp.status_code in (401, 403):
            raise KybernisProtocolException("unauthorized")

    def _report_task_completion(
        self,
        method_name: str,
        framework: str,
        task_id: str,
        run_id: str,
        result: Any,
        params: Dict[str, Any],
    ) -> None:
        if not self.gateway_url:
            return
        payload = {
            "jsonrpc": "2.0",
            "id": str(uuid.uuid4()),
            "method": "tasks/execute",
            "params": {
                "task_id": task_id,
                "framework": framework,
                "method": method_name,
                "run_id": run_id,
                "tenant_id": self.tenant_id,
                "input": self._safe_json(params),
                "output": self._safe_json(result),
            },
        }
        _ = self._post_a2a(payload)

    def _post_a2a(self, payload: Dict[str, Any]) -> Optional[requests.Response]:
        url = f"{self.gateway_url}/a2a/rpc"
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        if self.tenant_id:
            headers["X-Kyb-Tenant-Id"] = self.tenant_id
        run_id = payload.get("params", {}).get("run_id")
        if run_id:
            headers["X-Kyb-Run-Id"] = run_id
        try:
            return requests.post(url, headers=headers, data=json.dumps(payload), timeout=self.timeout_s)
        except Exception:
            return None

    def _post_mcp(self, payload: Dict[str, Any]) -> Optional[requests.Response]:
        url = f"{self.gateway_url}/mcp/invoke"
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        if self.tenant_id:
            headers["X-Kyb-Tenant-Id"] = self.tenant_id
        run_id = payload.get("run_id") or payload.get("input", {}).get("run_id")
        if isinstance(run_id, str) and run_id:
            headers["X-Kyb-Run-Id"] = run_id
        try:
            return requests.post(url, headers=headers, data=json.dumps(payload), timeout=self.timeout_s)
        except Exception:
            return None

    def _ensure_run_id(self, framework: str, obj: Any, args: tuple, kwargs: Dict[str, Any]) -> str:
        if framework == "langgraph":
            return self._get_or_create_run_id(kwargs, ["config", "configurable", "thread_id"])
        if framework == "crewai":
            return self._get_or_create_run_id(kwargs, ["inputs", "_kybernis_run_id"], prefix="kyb_crew_")
        if framework == "autogen":
            message = kwargs.get("message")
            if isinstance(message, dict):
                return self._get_or_create_run_id(kwargs, ["message", "kybernis_run_id"], prefix="kyb_ag_")
            chat_id = getattr(obj, "chat_id", "")
            if isinstance(chat_id, str) and chat_id:
                return chat_id
            return f"kyb_ag_{uuid.uuid4().hex[:8]}"
        if framework == "google_adk":
            rid = kyb_run_id.get()
            return rid or ""
        extracted = self._extract_run_id_from_kwargs(kwargs)
        if extracted:
            return extracted
        extracted = self._extract_run_id_from_obj(obj, args)
        if extracted:
            return extracted
        return kyb_run_id.get() or ""

    def _resolve_tool_name(self, method_name: str, obj: Any, args: tuple, kwargs: Dict[str, Any]) -> str:
        if self.tool_resolver is not None:
            try:
                resolved = self.tool_resolver(method_name, obj, args, kwargs)
                if isinstance(resolved, str):
                    return resolved
            except Exception:
                return ""
        return self.tool_map.get(method_name, "")

    def _maybe_invoke_mcp(
        self,
        method_name: str,
        framework: str,
        obj: Any,
        args: tuple,
        kwargs: Dict[str, Any],
        run_id: str,
    ) -> None:
        if not self.mode_mcp or not self.gateway_url:
            return
        tool_name = self._resolve_tool_name(method_name, obj, args, kwargs)
        if not tool_name:
            return
        payload = {
            "tool": tool_name,
            "tenant_id": self.tenant_id,
            "run_id": run_id,
            "input": {
                "framework": framework,
                "method": method_name,
                "args": self._safe_json(args),
                "kwargs": self._safe_json(kwargs),
            },
        }
        _ = self._post_mcp(payload)

    @staticmethod
    def _split_csv(value: str) -> Set[str]:
        parts = [p.strip() for p in value.split(",")]
        return {p for p in parts if p}

    def _extract_run_id_from_kwargs(self, kwargs: Dict[str, Any]) -> str:
        for key in ("run_id", "trace_id", "session_id", "thread_id", "execution_id"):
            val = kwargs.get(key)
            if isinstance(val, str) and val:
                return val
        for container_key in ("config", "context", "request", "message", "metadata"):
            container = kwargs.get(container_key)
            extracted = self._extract_run_id_from_container(container)
            if extracted:
                return extracted
        return ""

    def _extract_run_id_from_obj(self, obj: Any, args: tuple) -> str:
        for key in ("run_id", "trace_id", "session_id", "thread_id", "execution_id"):
            val = getattr(obj, key, "")
            if isinstance(val, str) and val:
                return val
        for arg in args:
            extracted = self._extract_run_id_from_container(arg)
            if extracted:
                return extracted
        return ""

    def _extract_run_id_from_container(self, container: Any) -> str:
        if container is None:
            return ""
        if isinstance(container, dict):
            for key in ("run_id", "trace_id", "session_id", "thread_id", "execution_id"):
                val = container.get(key)
                if isinstance(val, str) and val:
                    return val
            nested = container.get("session") or container.get("context") or container.get("metadata")
            return self._extract_run_id_from_container(nested)
        for key in ("run_id", "trace_id", "session_id", "thread_id", "execution_id"):
            val = getattr(container, key, "")
            if isinstance(val, str) and val:
                return val
        nested = getattr(container, "session", None) or getattr(container, "context", None)
        return self._extract_run_id_from_container(nested)

    def _get_or_create_run_id(self, kwargs: Dict[str, Any], path: list, prefix: str = "kyb_") -> str:
        current: Dict[str, Any] = kwargs
        for step in path[:-1]:
            if step not in current or not isinstance(current[step], dict):
                current[step] = {}
            current = current[step]
        last_key = path[-1]
        val = current.get(last_key)
        if not isinstance(val, str) or not val:
            val = f"{prefix}{uuid.uuid4().hex[:12]}"
            current[last_key] = val
        return val

    @staticmethod
    def _detect_framework(instance: Any) -> str:
        module = instance.__class__.__module__
        print(f"Detected framework module: {module}")
        print("google" in module or "vertexai" in module or "adk" in module)
        if "langgraph" in module:
            return "langgraph"
        if "crewai" in module:
            return "crewai"
        if "autogen" in module:
            return "autogen"
        if "google" in module or "vertexai" in module or "adk" in module:
            return "google_adk"
        return "generic"

    def _extract_adk_session_id(self, obj: Any, args: tuple, kwargs: Dict[str, Any]) -> str:
        session_path = kwargs.get("session_path")
        if isinstance(session_path, str) and "/" in session_path:
            return session_path.rstrip("/").split("/")[-1]
        session_id = kwargs.get("session_id")
        if isinstance(session_id, str) and session_id:
            return session_id
        request = kwargs.get("request")
        extracted = self._extract_session_from_container(request)
        if extracted:
            return extracted
        for arg in args:
            extracted = self._extract_session_from_container(arg)
            if extracted:
                return extracted
        obj_session = getattr(obj, "session_id", "")
        if isinstance(obj_session, str) and obj_session:
            return obj_session
        session_obj = getattr(obj, "session", None)
        extracted = self._extract_session_from_container(session_obj)
        if extracted:
            return extracted
        return ""

    @staticmethod
    def _extract_session_from_container(container: Any) -> str:
        if container is None:
            return ""
        if isinstance(container, dict):
            session_path = container.get("session_path")
            if isinstance(session_path, str) and "/" in session_path:
                return session_path.rstrip("/").split("/")[-1]
            session_id = container.get("session_id") or container.get("id")
            if isinstance(session_id, str) and session_id:
                return session_id
            session = container.get("session")
            if isinstance(session, dict):
                session_id = session.get("id") or session.get("session_id")
                if isinstance(session_id, str) and session_id:
                    return session_id
        session_path = getattr(container, "session_path", "")
        if isinstance(session_path, str) and "/" in session_path:
            return session_path.rstrip("/").split("/")[-1]
        session_id = getattr(container, "session_id", "")
        if isinstance(session_id, str) and session_id:
            return session_id
        session = getattr(container, "session", None)
        if session is not None:
            session_id = getattr(session, "id", "")
            if isinstance(session_id, str) and session_id:
                return session_id
        return ""

    @staticmethod
    def _safe_json(value: Any) -> Any:
        try:
            json.dumps(value)
            return value
        except Exception:
            return json.loads(json.dumps(value, default=str))


class _SDKProxy:
    def __init__(self, sdk: KybernisSDK, target: Any, framework: str) -> None:
        self._sdk = sdk
        self._target = target
        self._framework = framework

    def __getattr__(self, name: str) -> Any:
        if name.startswith("_"):
            return getattr(self._target, name)
        if self._sdk.allow_methods and name not in self._sdk.allow_methods:
            return getattr(self._target, name)
        if name in self._sdk.exclude_methods:
            return getattr(self._target, name)
        attr = getattr(self._target, name)
        if inspect.iscoroutinefunction(attr):
            @functools.wraps(attr)
            async def async_wrapper(*args, **kwargs):
                task_id = kwargs.get("task_id") or f"task_{uuid.uuid4().hex}"
                run_id = self._sdk._ensure_run_id(self._framework, self._target, args, kwargs)
                await self._sdk._run_async(self._sdk._check_protocol_compliance, name, self._framework, task_id, run_id)
                await self._sdk._run_async(self._sdk._maybe_invoke_mcp, name, self._framework, self._target, args, kwargs, run_id)
                result = await attr(*args, **kwargs)
                await self._sdk._run_async(self._sdk._report_task_completion, name, self._framework, task_id, run_id, result, kwargs)
                return result
            return async_wrapper
        if callable(attr):
            @functools.wraps(attr)
            def wrapper(*args, **kwargs):
                task_id = kwargs.get("task_id") or f"task_{uuid.uuid4().hex}"
                run_id = self._sdk._ensure_run_id(self._framework, self._target, args, kwargs)
                self._sdk._check_protocol_compliance(name, self._framework, task_id, run_id)
                self._sdk._maybe_invoke_mcp(name, self._framework, self._target, args, kwargs, run_id)
                result = attr(*args, **kwargs)
                self._sdk._report_task_completion(name, self._framework, task_id, run_id, result, kwargs)
                return result
            return wrapper
        return attr
