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

import asyncio
import functools
import json
import os
import time
import uuid
from typing import Any, AsyncGenerator, Callable, Optional

try:
    from google.adk.agents.base_agent import BaseAgent  # type: ignore
    from google.adk.tools.base_tool import BaseTool     # type: ignore
    from google.adk.tools import ToolContext            # type: ignore
except Exception:  # pragma: no cover
    BaseAgent = None  # type: ignore
    BaseTool = None   # type: ignore
    ToolContext = None  # type: ignore


def _get_invocation_id_from_agent_parent_ctx(parent_ctx: Any) -> str:
    """
    ADK guarantees InvocationContext has invocation_id.
    BaseAgent.run_async creates ctx = parent_context.model_copy(update={'agent': self}).
    So parent_ctx.invocation_id is the correct run_id.
    """
    inv = getattr(parent_ctx, "invocation_id", None)
    if isinstance(inv, str) and inv:
        return inv
    # Fallback for safety (should almost never happen)
    return f"kyb_adk_{uuid.uuid4().hex[:12]}"


def _get_invocation_id_from_tool_context(tool_context: Any) -> str:
    inv_ctx = getattr(tool_context, "invocation_context", None)
    inv = getattr(inv_ctx, "invocation_id", None)
    if isinstance(inv, str) and inv:
        return inv
    return f"kyb_adk_{uuid.uuid4().hex[:12]}"


def _temporary_env(env: dict[str, str]):
    class _Tmp:
        def __enter__(self):
            self._old = {}
            for k, v in env.items():
                self._old[k] = os.environ.get(k)
                os.environ[k] = v
            return self
        def __exit__(self, exc_type, exc, tb):
            for k, old in self._old.items():
                if old is None:
                    os.environ.pop(k, None)
                else:
                    os.environ[k] = old
            return False
    return _Tmp()


class KybernisAdkPatcher:
    """
    Global one-time patcher for ADK:
      - intercept BaseAgent.run_async / run_live to establish run_id
      - intercept BaseTool.run_async to enforce receipts, retry/budget policies
    """

    _installed = False

    def __init__(self, sdk: Any):
        self.sdk = sdk  # KybernisSDK instance (has _check_protocol_compliance, etc.)

    def install_once(self) -> None:
        if KybernisAdkPatcher._installed:
            return
        if BaseAgent is None or BaseTool is None:
            return
        self._patch_base_agent()
        self._patch_base_tool()
        KybernisAdkPatcher._installed = True

    def _patch_base_agent(self) -> None:
        original_run_async = BaseAgent.run_async
        original_run_live = BaseAgent.run_live

        @functools.wraps(original_run_async)
        async def run_async_wrapped(agent_self: Any, parent_context: Any) -> AsyncGenerator[Any, None]:
            run_id = _get_invocation_id_from_agent_parent_ctx(parent_context)
            task_id = f"task_{uuid.uuid4().hex}"

            # Optional: tell Kybernis “a run started”
            await self._maybe_call_a2a_stream(method="run_async", task_id=task_id, run_id=run_id)

            with _temporary_env({"KYB_RUN_ID": run_id, "KYB_SESSION_ID": run_id}):
                # yield through the original generator
                async for event in original_run_async(agent_self, parent_context):
                    yield event

            # Optional: tell Kybernis “run finished”
            await self._maybe_call_a2a_execute(method="run_async", task_id=task_id, run_id=run_id, result={"status": "ok"})

        @functools.wraps(original_run_live)
        async def run_live_wrapped(agent_self: Any, parent_context: Any) -> AsyncGenerator[Any, None]:
            run_id = _get_invocation_id_from_agent_parent_ctx(parent_context)
            task_id = f"task_{uuid.uuid4().hex}"
            await self._maybe_call_a2a_stream(method="run_live", task_id=task_id, run_id=run_id)

            with _temporary_env({"KYB_RUN_ID": run_id, "KYB_SESSION_ID": run_id}):
                async for event in original_run_live(agent_self, parent_context):
                    yield event

            await self._maybe_call_a2a_execute(method="run_live", task_id=task_id, run_id=run_id, result={"status": "ok"})

        # patch the class (global)
        BaseAgent.run_async = run_async_wrapped  # type: ignore
        BaseAgent.run_live = run_live_wrapped    # type: ignore

    def _patch_base_tool(self) -> None:
        original_tool_run_async = BaseTool.run_async

        @functools.wraps(original_tool_run_async)
        async def tool_run_async_wrapped(tool_self: Any, *, args: dict[str, Any], tool_context: Any) -> Any:
            run_id = _get_invocation_id_from_tool_context(tool_context)
            tool_name = getattr(tool_self, "name", None) or tool_self.__class__.__name__

            # Receipt start (a2a/mcp/proxy — your choice)
            receipt_id = f"rcpt_{uuid.uuid4().hex[:10]}"
            start_t = time.time()
            await self._maybe_emit_receipt(
                phase="start",
                run_id=run_id,
                receipt_id=receipt_id,
                tool_name=tool_name,
                payload={"args": self._safe_json(args)},
            )

            # Governance point: allow/deny/rewrite args BEFORE side-effect
            # (If you want allowlists, budgets, approvals: do it here.)
            # Example hook:
            # decision = self.sdk.policy.decide(run_id, tool_name, args, tool_context)
            # if decision.block: raise KybernisProtocolException("blocked")

            try:
                result = await original_tool_run_async(tool_self, args=args, tool_context=tool_context)
                ok = True
                return result
            except Exception as e:
                ok = False
                # Optional: deterministic retries (bounded)
                # If you implement retries, do it here with strict caps.
                raise
            finally:
                dur_ms = int((time.time() - start_t) * 1000)
                await self._maybe_emit_receipt(
                    phase="end",
                    run_id=run_id,
                    receipt_id=receipt_id,
                    tool_name=tool_name,
                    payload={"ok": ok, "duration_ms": dur_ms},
                )

        BaseTool.run_async = tool_run_async_wrapped  # type: ignore

    async def _maybe_call_a2a_stream(self, method: str, task_id: str, run_id: str) -> None:
        # Reuse your existing sync method by running it in a thread, or add async requests.
        await self._to_thread(self.sdk._check_protocol_compliance, method, "google_adk", task_id, run_id)

    async def _maybe_call_a2a_execute(self, method: str, task_id: str, run_id: str, result: Any) -> None:
        await self._to_thread(self.sdk._report_task_completion, method, "google_adk", task_id, run_id, result, {})

    async def _maybe_emit_receipt(self, phase: str, run_id: str, receipt_id: str, tool_name: str, payload: dict[str, Any]) -> None:
        # Keep this extremely small + fast for demos.
        # You can route to your gateway_url if present.
        if not getattr(self.sdk, "gateway_url", ""):
            return
        body = {
            "run_id": run_id,
            "receipt_id": receipt_id,
            "phase": phase,
            "tool": tool_name,
            "payload": payload,
        }
        await self._to_thread(self.sdk._post_mcp, {"tool": "kyb.receipt", "run_id": run_id, "input": body})

    async def _to_thread(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
        try:
            return await asyncio.to_thread(fn, *args, **kwargs)
        except AttributeError:
            loop = asyncio.get_running_loop()
            return await loop.run_in_executor(None, functools.partial(fn, *args, **kwargs))

    @staticmethod
    def _safe_json(value: Any) -> Any:
        try:
            json.dumps(value)
            return value
        except Exception:
            return json.loads(json.dumps(value, default=str))