"""
Agno Agent Integration Example — infralo-sdk

Demonstrates:
  1. Wrapping an Agno Agent with Infralo tool + model callbacks.
  2. Full tool span capture with inputs, outputs, and status.
  3. LLM span linking via the Infralo gateway (x-infralo-span-id header).

Requirements:
    pip install infralo[agno]
    # Set env vars or edit constants below.

Run:
    cd examples/agno
    uv run agno_integration.py
    # or: python agno_integration.py
"""

import asyncio
import json
import logging
import os

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

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

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 (replace with your real tools) ──────────────────────────────


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


async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    logger.info("  [tool] get_weather(city=%r)", city)
    return json.dumps({"city": city, "temp_c": 18, "condition": "partly cloudy"})


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


async def run_agno_example() -> None:
    try:
        from agno.agent import Agent  # type: ignore[import-untyped]
        from agno.models.openai import OpenAIChat  # type: ignore[import-untyped]
    except ImportError:
        logger.error("agno is not installed. Run: pip install infralo[agno]")
        return

    print("=== Agno + Infralo Integration Example ===")

    # Build the callback hooks — one call wires up all four lifecycle hooks.
    callbacks = make_agno_callbacks(infralo)

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

        agent = Agent(
            # Route through the Infralo gateway for LLM span linkage
            model=OpenAIChat(
                id=MODEL_DEPLOYMENT,
                base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
                api_key=INFRALO_API_KEY,
            ),
            tools=[search_web, get_weather],
            instructions=["You are a helpful assistant that uses tools to answer questions."],
            markdown=False,
            # Inject all Infralo lifecycle hooks
            **callbacks,
        )

        response = await agent.arun(
            "What is the weather in Tokyo and what's trending about AI today?"
        )
        print(f"\nAgent response:\n{getattr(response, 'content', str(response))}")

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


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