"""
Agno 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/agno
    uv run agno_complex_nesting.py
"""

import asyncio
import json
import logging
import os

from infralo import Infralo
from infralo.integrations.agno import make_agno_callbacks

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

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")

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

# Shared lifecycle callbacks for Agno
callbacks = make_agno_callbacks(infralo)


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


async def search_web(query: str) -> str:
    """Search the web for up-to-date information."""
    logger.info("    [Nested Tool] search_web(query=%r)", query)
    return json.dumps(
        {
            "results": [
                "LangGraph & Agno are trending frameworks in 2026 for building agentic architectures.",
                f"Found news result for '{query}': Multi-agent frameworks are rapidly evolving.",
            ]
        }
    )


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


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

    from agno.agent import Agent
    from agno.models.openai import OpenAIChat

    # Create the specialized Research Agent
    research_agent = Agent(
        name="ResearchAgent",
        model=OpenAIChat(
            id=MODEL_DEPLOYMENT,
            base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
            api_key=INFRALO_API_KEY,
        ),
        tools=[search_web],
        instructions=[
            "You are a research specialist. Search the web to compile insights on the given topic.",
            "Return a concise summary of your findings.",
        ],
        markdown=False,
        # Register the callbacks so this agent's actions are traced under the active trace context
        **callbacks,
    )

    # Run the nested agent. The global HTTPX patch automatically carries the active parent span!
    response = await research_agent.arun(f"Research the following topic: {topic}")
    content = getattr(response, "content", str(response))

    logger.info("  [Root Tool] delegation completed. Summary of results returned.")
    return content


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


async def run_complex_nesting():
    try:
        from agno.agent import Agent
        from agno.models.openai import OpenAIChat
    except ImportError:
        logger.error("agno is not installed. Run: pip install infralo[agno]")
        return

    print("=== Running Agno Complex Nesting Example ===")

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

        # Create the Root Agent
        root_agent = Agent(
            name="RootAgent",
            model=OpenAIChat(
                id=MODEL_DEPLOYMENT,
                base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
                api_key=INFRALO_API_KEY,
            ),
            tools=[delegate_research],
            instructions=[
                "You are an orchestrator assistant.",
                "To answer research queries, delegate them to the ResearchAgent tool.",
            ],
            markdown=False,
            # Register the callbacks for the root agent
            **callbacks,
        )

        response = await root_agent.arun("Can you perform deep research on multi-agent frameworks?")
        print(f"\nFinal Root Agent response:\n{getattr(response, 'content', str(response))}")

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


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