# kybernis/sdk_adk_instance_patch.py
from __future__ import annotations

import uuid
import functools
from typing import Any, AsyncGenerator

from .context import kyb_run_id, kyb_session_id

def _new_run_id() -> str:
    return f"kyb_{uuid.uuid4().hex[:12]}"

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

        # --- run_async ---
        if hasattr(agent, "run_async"):
            original_bound = agent.run_async  # bound method (expects parent_context)

            @functools.wraps(original_bound)
            async def run_async_wrapped(parent_context: Any) -> AsyncGenerator[Any, None]:
                # Prefer session.id as run_id (UUID). Fallback to invocation_id.
                inv_ctx = getattr(parent_context, "invocation_context", None)
                if inv_ctx is not None:
                    inv_id = getattr(inv_ctx, "invocation_id", None)
                    sess_obj = getattr(inv_ctx, "session", None)
                    sess_id = getattr(sess_obj, "id", None) if sess_obj is not None else None
                else:
                    inv_id = getattr(parent_context, "invocation_id", None)
                    sess_obj = getattr(parent_context, "session", None)
                    sess_id = getattr(sess_obj, "id", None) if sess_obj is not None else None

                run_id = sess_id or inv_id or _new_run_id()
                session_id = sess_id or run_id

                prev_run = kyb_run_id.get()
                created_here = prev_run is None

                token_run = kyb_run_id.set(run_id) if created_here else None
                token_sess = kyb_session_id.set(session_id) if created_here else None

                try:
                    async for ev in original_bound(parent_context):
                        yield ev
                finally:
                    if created_here:
                        kyb_run_id.reset(token_run)   # type: ignore[arg-type]
                        kyb_session_id.reset(token_sess)  # type: ignore[arg-type]

            # IMPORTANT: assign directly (no MethodType) to avoid extra self arg
            try:
                agent.run_async = run_async_wrapped
            except Exception:
                object.__setattr__(agent, "run_async", run_async_wrapped)

        # --- run_live ---
        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_ctx = getattr(parent_context, "invocation_context", None)
                if inv_ctx is not None:
                    inv_id = getattr(inv_ctx, "invocation_id", None)
                    sess_obj = getattr(inv_ctx, "session", None)
                    sess_id = getattr(sess_obj, "id", None) if sess_obj is not None else None
                else:
                    inv_id = getattr(parent_context, "invocation_id", None)
                    sess_obj = getattr(parent_context, "session", None)
                    sess_id = getattr(sess_obj, "id", None) if sess_obj is not None else None

                run_id = sess_id or inv_id or _new_run_id()
                session_id = sess_id or run_id

                prev_run = kyb_run_id.get()
                created_here = prev_run is None

                token_run = kyb_run_id.set(run_id) if created_here else None
                token_sess = kyb_session_id.set(session_id) if created_here else None

                try:
                    async for ev in original_live_bound(parent_context):
                        yield ev
                finally:
                    if created_here:
                        kyb_run_id.reset(token_run)   # type: ignore[arg-type]
                        kyb_session_id.reset(token_sess)  # type: ignore[arg-type]

            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)
