from typing import Any, Callable, Optional, Tuple

from ._util import temporary_env
from .base import FrameworkAdapter, adapter_registry

try:
    import beeai  # type: ignore
    _BEEAI_AVAILABLE = True
except Exception:
    _BEEAI_AVAILABLE = False


def _extract_context(args: Tuple[Any, ...], kwargs: dict) -> Tuple[Optional[object], Tuple[Any, ...], dict]:
    if "context" in kwargs:
        ctx = kwargs.pop("context")
        return ctx, args, kwargs
    if "state" in kwargs:
        ctx = kwargs.get("state")
        return ctx, args, kwargs
    if args and isinstance(args[-1], dict):
        return args[-1], args[:-1], kwargs
    return None, args, kwargs


def _run_id_from_context(ctx: Optional[object]) -> str:
    if not ctx:
        return ""
    if isinstance(ctx, dict):
        for key in ("run_id", "trace_id", "session_id", "conversation_id"):
            val = ctx.get(key)
            if isinstance(val, str) and val:
                return val
    for key in ("run_id", "trace_id", "session_id", "conversation_id"):
        val = getattr(ctx, key, "")
        if isinstance(val, str) and val:
            return val
    return ""


def wrap_beeai_tool(fn: Callable[..., object]) -> Callable[..., object]:
    def wrapper(*args: Any, **kwargs: Any) -> object:
        ctx, new_args, new_kwargs = _extract_context(args, kwargs)
        run_id = _run_id_from_context(ctx)
        env = {"KYB_RUN_ID": run_id} if run_id else {}
        with temporary_env(env):
            return fn(*new_args, **new_kwargs)

    return wrapper


def _beeai_available() -> bool:
    return _BEEAI_AVAILABLE


beeai_adapter = FrameworkAdapter("beeai", _beeai_available, wrap_beeai_tool)
adapter_registry.register(beeai_adapter)
