# kybernis/adapters/adk.py
from __future__ import annotations

import asyncio
import functools
from typing import Any, Callable, Optional, Tuple

from .base import FrameworkAdapter, adapter_registry

try:
    from google.adk.tools import ToolContext  # type: ignore
    _ADK_AVAILABLE = True
except Exception:
    ToolContext = None  # type: ignore
    _ADK_AVAILABLE = False

from ..context import kyb_run_id
from ..runtime import kyb_client, resolve_run_id_from_invocation_id
from .._asyncutil import fire_and_forget
from ..spans import span as kyb_span


def _extract_tool_context(args: Tuple[Any, ...], kwargs: dict) -> Tuple[Optional[object], Tuple[Any, ...], dict]:
    if "tool_context" in kwargs:
        ctx = kwargs.get("tool_context")
        return ctx, args, kwargs
    if ToolContext is not None and args and isinstance(args[-1], ToolContext):
        return args[-1], args[:-1], kwargs
    return None, args, kwargs


def _run_ids_from_tool_context(tool_context: Optional[object]) -> tuple[str, str, str]:
    if not tool_context:
        return "", "", ""
    inv = getattr(tool_context, "invocation_context", None)
    inv_id = str(getattr(inv, "invocation_id", "") or "") if inv else ""
    sess = getattr(inv, "session", None) if inv else None
    sess_id = str(getattr(sess, "id", "") or "") if sess else ""
    call_id = str(getattr(tool_context, "function_call_id", "") or "")
    return inv_id, sess_id, call_id


def _tool_display_name(fn: Any) -> str:
    return getattr(fn, "_name", "") or getattr(fn, "name", "") or fn.__name__ if hasattr(fn, "__name__") else str(fn)


def wrap_adk_tool(fn: Callable[..., Any]) -> Callable[..., Any]:
    """
    Wraps a callable tool for ADK.
    - No env vars.
    - Sets ContextVars from tool_context if present.
    - Emits a TOOL span around the call.
    - Does not block on Kybernis network (fire-and-forget) unless you explicitly choose to.
    """
    if getattr(fn, "_kyb_wrapped", False):
        return fn

    is_async = asyncio.iscoroutinefunction(fn)

    if is_async:
        @functools.wraps(fn)
        async def awrapper(*args: Any, **kwargs: Any) -> Any:
            tool_context, args2, kwargs2 = _extract_tool_context(args, kwargs)

            inv_id, _sess_id, call_id = _run_ids_from_tool_context(tool_context)

            tok_run = None
            run_id = ""

            if inv_id:
                # Map ADK invocation_id -> Kybernis UUID run_id (and set kyb_run_id)
                run_id = resolve_run_id_from_invocation_id(inv_id)
                tok_run = kyb_run_id.set(run_id)  # get a reset token for hygiene
            else:
                # Nested tool calls may not have tool_context; reuse current run_id if any
                existing = kyb_run_id.get()
                if existing:
                    run_id = existing
                    tok_run = kyb_run_id.set(existing)

            name = _tool_display_name(fn)
            attrs = {"tool": {"tool_name": name}}
            if call_id:
                attrs["tool"]["function_call_id"] = call_id

            client = kyb_client()

            # fire-and-forget RUN_START-like event (optional)
            if run_id:
                fire_and_forget(lambda: client.emit_event(run_id, "TOOL_SEEN", {"tool": name, "function_call_id": call_id}))

            try:
                if run_id:
                    # span emits to kybernis; keep it lightweight
                    with kyb_span(client=client, span_type="tool", name=name, attributes=attrs):
                        return await fn(*args2, **kwargs2)
                return await fn(*args2, **kwargs2)
            finally:
                if tok_run is not None:
                    kyb_run_id.reset(tok_run)

        setattr(awrapper, "_kyb_wrapped", True)
        return awrapper

    @functools.wraps(fn)
    def swrapper(*args: Any, **kwargs: Any) -> Any:
        tool_context, args2, kwargs2 = _extract_tool_context(args, kwargs)

        inv_id, _sess_id, call_id = _run_ids_from_tool_context(tool_context)

        tok_run = None
        run_id = ""

        if inv_id:
            run_id = resolve_run_id_from_invocation_id(inv_id)
            tok_run = kyb_run_id.set(run_id)
        else:
            existing = kyb_run_id.get()
            if existing:
                run_id = existing
                tok_run = kyb_run_id.set(existing)

        name = _tool_display_name(fn)
        attrs = {"tool": {"tool_name": name}}
        if call_id:
            attrs["tool"]["function_call_id"] = call_id

        client = kyb_client()
        if run_id:
            fire_and_forget(lambda: client.emit_event(run_id, "TOOL_SEEN", {"tool": name, "function_call_id": call_id}))

        try:
            if run_id:
                with kyb_span(client=client, span_type="tool", name=name, attributes=attrs):
                    return fn(*args2, **kwargs2)
            return fn(*args2, **kwargs2)
        finally:
            if tok_run is not None:
                kyb_run_id.reset(tok_run)

    setattr(swrapper, "_kyb_wrapped", True)
    return swrapper


def _adk_available() -> bool:
    return _ADK_AVAILABLE


adk_adapter = FrameworkAdapter("adk", _adk_available, wrap_adk_tool)
adapter_registry.register(adk_adapter)
