"""
Google ADK Complex Nesting Example — infralo-sdk

Demonstrates:
  1. Multi-agent delegation: a Root Agent delegating to a Research Agent.
  2. Nested LLM and tool trace parenting.
  3. Automatic context propagation across nested agent boundaries.

Run:
    cd examples/google_adk
    uv run python google_adk_complex_nesting.py
"""

import asyncio
import logging
import os

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

# Configure logging
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", "vk_AmEFYr-k-S3_PN_i91wIIM19yVfV5PUj7T4sTdCeOSQ"
)
INFRALO_ENDPOINT = os.environ.get("INFRALO_ENDPOINT", "http://localhost:8000")
MODEL_DEPLOYMENT = os.environ.get("INFRALO_MODEL_DEPLOYMENT", "your_model_name")

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

# Shared lifecycle callbacks for Google ADK
callbacks = make_google_adk_callbacks(infralo)


# ── Nested Tool ──────────────────────────────────────────────────────────────


async def search_web(query: str) -> dict:
    """Search the web for up-to-date information."""
    logger.info("    [Nested Tool] search_web(query=%r)", query)
    return {
        "results": [
            "Google ADK is a powerful framework for building agentic workflows.",
            f"Found news result for '{query}': Multi-agent frameworks are rapidly evolving in 2026.",
        ]
    }


# ── Root Agent Tool (Delegates to Nested Agent) ──────────────────────────────


async def delegate_research(topic: str) -> dict:
    """Delegate in-depth research about a specific topic to a specialized Research Agent."""
    logger.info("  [Root Tool] delegating research for topic=%r...", topic)

    from google.adk.agents import Agent
    from google.adk.models.lite_llm import LiteLlm
    from google.adk.runners import Runner
    from google.adk.sessions import InMemorySessionService
    from google.genai.types import Content, Part

    # Create the specialized Research Agent
    research_agent = Agent(
        name="ResearchAgent",
        model=LiteLlm(
            model=f"openai/{MODEL_DEPLOYMENT}",
            api_base=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
            api_key=INFRALO_API_KEY,
        ),
        tools=[search_web],
        instruction="You are a research specialist. Search the web to compile insights on the given topic and return a concise summary.",
        **callbacks,
    )

    session_service = InMemorySessionService()
    runner = Runner(
        agent=research_agent, app_name="research-subsystem", session_service=session_service
    )
    session = await session_service.create_session(
        app_name="research-subsystem", user_id="demo-user"
    )

    user_message = Content(
        parts=[Part(text=f"Research the following topic: {topic}")],
        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

    summary = (
        final_response.content.parts[0].text if final_response and final_response.content else ""
    )
    logger.info("  [Root Tool] delegation completed. Summary of results returned.")
    return {"summary": summary}


# ── Main Execution ────────────────────────────────────────────────────────────


async def run_complex_nesting():
    try:
        from google.adk.agents import Agent
        from google.adk.models.lite_llm import LiteLlm
        from google.adk.runners import Runner
        from google.adk.sessions import InMemorySessionService
        from google.genai.types import Content, Part
    except ImportError:
        logger.error("google-adk is not installed. Run: pip install infralo[google-adk]")
        return

    print("=== Running Google ADK Complex Nesting Example ===")

    session_service = InMemorySessionService()

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

        # Create the Root Agent
        root_agent = Agent(
            name="RootAgent",
            model=LiteLlm(
                model=f"openai/{MODEL_DEPLOYMENT}",
                api_base=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
                api_key=INFRALO_API_KEY,
            ),
            tools=[delegate_research],
            instruction="You are an orchestrator assistant. To answer research queries, delegate them to the delegate_research tool.",
            **callbacks,
        )

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

        user_message = Content(
            parts=[Part(text="Can you perform deep research on multi-agent frameworks?")],
            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"\nFinal Root Agent response:\n{final_response.content.parts[0].text}")

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


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