"""
Google ADK Integration Example — infralo-sdk

Demonstrates:
  1. Using make_google_adk_callbacks() to inject Infralo hooks into a Google ADK Agent.
  2. Full tool span capture with inputs, outputs, and error status.
  3. LLM span linking via the Infralo gateway.

Requirements:
    pip install infralo[google-adk]

Run:
    cd examples/google_adk
    uv run google_adk_integration.py
    # or: python google_adk_integration.py
"""

import asyncio
import logging
import os

from infralo import Infralo
from infralo.integrations.google_adk import make_google_adk_callbacks

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("infralo.examples.google_adk")

INFRALO_API_KEY = os.environ.get(
    "INFRALO_API_KEY", "your_workspace_api_key"
)
INFRALO_ENDPOINT = os.environ.get("INFRALO_ENDPOINT", "http://localhost:8000")
MODEL_DEPLOYMENT = os.environ.get("INFRALO_MODEL_DEPLOYMENT", "your_model_name")

# ── 1. Infralo client ─────────────────────────────────────────────────────────

infralo = Infralo(
    api_key=INFRALO_API_KEY,
    endpoint=INFRALO_ENDPOINT,
    flush_interval=1.0,
)

# ── 2. Stub tools ─────────────────────────────────────────────────────────────


async def search_web(query: str) -> dict:
    """Search the web for up-to-date information on a topic."""
    logger.info("  [tool] search_web(query=%r)", query)
    return {"results": [f"Result for '{query}': item 1", f"Result for '{query}': item 2"]}


async def get_weather(city: str) -> dict:
    """Get the current weather for a city."""
    logger.info("  [tool] get_weather(city=%r)", city)
    return {"city": city, "temperature_c": 20, "condition": "sunny"}


# ── 3. Agent run ──────────────────────────────────────────────────────────────


async def run_google_adk_example() -> None:
    try:
        from google.adk.agents import Agent  # type: ignore[import-untyped]
        from google.adk.models.lite_llm import LiteLlm  # type: ignore[import-untyped]
        from google.adk.runners import Runner  # type: ignore[import-untyped]
        from google.adk.sessions import InMemorySessionService  # type: ignore[import-untyped]
    except ImportError:
        logger.error("google-adk is not installed. Run: pip install infralo[google-adk]")
        return

    print("=== Google ADK + Infralo Integration Example ===")

    callbacks = make_google_adk_callbacks(infralo)

    agent = Agent(
        name="infralo_demo_agent",
        # Route through the Infralo gateway for full LLM+tool span linkage
        model=LiteLlm(
            model=f"openai/{MODEL_DEPLOYMENT}",
            api_base=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
            api_key=INFRALO_API_KEY,
        ),
        tools=[search_web, get_weather],
        instruction="You are a helpful assistant. Always use tools to answer questions.",
        **callbacks,
    )

    session_service = InMemorySessionService()

    with infralo.start_trace(session_id="google-adk-demo") as trace:
        print(f"Trace ID: {trace.trace_id}")

        runner = Runner(agent=agent, app_name="infralo-demo", session_service=session_service)
        session = await session_service.create_session(app_name="infralo-demo", user_id="demo-user")

        from google.genai.types import Content, Part  # type: ignore[import-untyped]

        user_message = Content(
            parts=[Part(text="What is the weather in Berlin? Also search for Google ADK news.")],
            role="user",
        )

        final_response = None
        async for event in runner.run_async(
            user_id="demo-user",
            session_id=session.id,
            new_message=user_message,
        ):
            if event.is_final_response():
                final_response = event

        if final_response and final_response.content:
            print(f"\nAgent response:\n{final_response.content.parts[0].text}")

    print("\nDone. Check Infralo dashboard for trace spans.")


if __name__ == "__main__":
    asyncio.run(run_google_adk_example())
    infralo.flush()