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

import requests

from .sdk_adk_instance_patch import patch_adk_agent_instance

from .runtime import kyb_client, ensure_run_id
from ._asyncutil import fire_and_forget
from .auto import install as install_http_patches
from .adapters._util import temporary_env
from .adapters import wrap_adk_tool

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:
    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,
    ) -> 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
        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

    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 = ensure_run_id()
            client = kyb_client()

            # pre: emit
            fire_and_forget(lambda: client.emit_event(run_id, "tool.call", {"tool": getattr(fn, "__name__", "tool")}))

            try:
                out = fn(*args, **kwargs)
                fire_and_forget(lambda: client.emit_event(run_id, "tool.result", {"tool": getattr(fn, "__name__", "tool"), "ok": True}))
                return out
            except Exception as e:
                fire_and_forget(lambda: client.emit_event(run_id, "tool.error", {"tool": getattr(fn, "__name__", "tool"), "error": str(e)}))
                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)
            run_id = getattr(inv_ctx, "invocation_id", None) or ensure_run_id()
            sess = getattr(inv_ctx, "session", None)
            session_id = getattr(sess, "id", "") or run_id

            # set contextvars (do not rely on env)
            token_run = kyb_run_id.set(run_id)
            token_sess = kyb_session_id.set(session_id)

            client = kyb_client()
            tool_name = getattr(tool, "name", tool.__class__.__name__)

            # pre: async, non-blocking
            fire_and_forget(lambda: client.emit_event(run_id, "tool.call", {"tool": tool_name, "args": self._safe_json(args)}))

            try:
                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}))
                return res
            except Exception as e:
                fire_and_forget(lambda: client.emit_event(run_id, "tool.error", {"tool": tool_name, "error": str(e)}))
                raise
            finally:
                kyb_run_id.reset(token_run)
                kyb_session_id.reset(token_sess)

        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_tool_if_needed(self, tool: Any) -> Any:
        # Case A: callable function tool
        if callable(tool) and not _is_adk_basetool(tool):
            if _mark(tool, "_kyb_wrapped"):
                return tool
            wrapped = wrap_adk_tool(tool)
            _mark(wrapped, "_kyb_wrapped")
            return wrapped

        # Case B: BaseTool instance → wrap its run_async
        if _is_adk_basetool(tool):
            if _mark(tool, "_kyb_run_async_wrapped"):
                return tool

            original = tool.run_async  # bound method

            @functools.wraps(original)
            async def run_async_wrapped(*, args: dict, tool_context: Any) -> Any:
                # This wrapper is called exactly like ADK calls it: keyword-only
                # It will set KYB_RUN_ID based on invocation_context.invocation_id
                print("Inside run_async_wrapped, tool_context:", tool_context)
                return await wrap_adk_tool(lambda **kw: original(**kw))(
                    args=args, tool_context=tool_context
                )

            try:
                tool.run_async = run_async_wrapped  # instance method override
            except Exception:
                object.__setattr__(tool, "run_async", run_async_wrapped)

            return tool

        # Case C: unknown / toolset / factories → leave as-is
        return tool

    def _wrap_adk_tools(self, agent: Any) -> None:
        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":
                    
                    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":
                    
                    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":
            extracted = self._extract_adk_session_id(obj, args, kwargs)
            if extracted:
                return extracted
            generated = f"kyb_adk_{uuid.uuid4().hex[:8]}"
            kwargs["session_id"] = generated
            return generated
        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 os.getenv("KYB_RUN_ID", "") or ensure_run_id() or uuid.uuid4().hex

    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
