# kybernis/sdk_adk_instance_patch.py
from __future__ import annotations

import functools
from typing import Any, AsyncGenerator

from .context import kyb_run_id
from .runtime import resolve_run_id_from_invocation_id

def _extract_invocation_id(parent_context: Any) -> str:
    inv_ctx = getattr(parent_context, "invocation_context", None)
    if inv_ctx is not None:
        inv_id = getattr(inv_ctx, "invocation_id", None)
    else:
        inv_id = getattr(parent_context, "invocation_id", None)
    inv_id = (str(inv_id or "")).strip()
    return inv_id

def patch_adk_agent_instance(root_agent: Any, *, also_patch_subagents: bool = True) -> None:
    def _patch_one(agent: Any) -> None:
        if getattr(agent, "_kyb_instance_patched", False):
            return

        if hasattr(agent, "run_async"):
            original_bound = agent.run_async

            @functools.wraps(original_bound)
            async def run_async_wrapped(parent_context: Any) -> AsyncGenerator[Any, None]:
                # Bind Kybernis run UUID at the start of the invocation (ASAP).
                inv_id = _extract_invocation_id(parent_context)
                tok_run = None
                if inv_id and kyb_run_id.get() is None:
                    run_id = resolve_run_id_from_invocation_id(inv_id)  # creates/loads run row
                    tok_run = kyb_run_id.set(run_id)

                try:
                    async for ev in original_bound(parent_context):
                        yield ev
                finally:
                    if tok_run is not None:
                        kyb_run_id.reset(tok_run)

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

        if hasattr(agent, "run_live"):
            original_live_bound = agent.run_live

            @functools.wraps(original_live_bound)
            async def run_live_wrapped(parent_context: Any) -> AsyncGenerator[Any, None]:
                inv_id = _extract_invocation_id(parent_context)
                tok_run = None
                if inv_id and kyb_run_id.get() is None:
                    run_id = resolve_run_id_from_invocation_id(inv_id)
                    tok_run = kyb_run_id.set(run_id)

                try:
                    async for ev in original_live_bound(parent_context):
                        yield ev
                finally:
                    if tok_run is not None:
                        kyb_run_id.reset(tok_run)

            try:
                agent.run_live = run_live_wrapped
            except Exception:
                object.__setattr__(agent, "run_live", run_live_wrapped)

        try:
            setattr(agent, "_kyb_instance_patched", True)
        except Exception:
            object.__setattr__(agent, "_kyb_instance_patched", True)

    _patch_one(root_agent)

    if also_patch_subagents:
        for sub in getattr(root_agent, "sub_agents", []) or []:
            patch_adk_agent_instance(sub, also_patch_subagents=True)