"""
LangGraph Complex Nesting Example — infralo-sdk

Demonstrates:
  1. Multi-graph delegation: a Parent Graph delegating to a Child Graph.
  2. Nested LLM and tool trace parenting in LangGraph.
  3. Automatic trace propagation across graph boundaries.

Run:
    cd examples/langgraph
    uv run langgraph_complex_nesting.py
"""

import asyncio
import logging
import os

from infralo import Infralo
from infralo.integrations.langgraph import InfraloCallbackHandler

# Configure logging
logging.basicConfig(
    level=logging.DEBUG, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("infralo.examples.langgraph_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,
)


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


def search_web(query: str) -> str:
    """Search the web for up-to-date information."""
    logger.info("    [Nested Tool] search_web(query=%r)", query)
    return f"Web results for '{query}': LangGraph subgraphs are fully traceable with Infralo."


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


async def delegate_research_tool(topic: str) -> str:
    """Call the specialized Child Agent to perform nested research."""
    from langchain_core.messages import HumanMessage
    from langchain_core.tools import tool as lc_tool
    from langchain_openai import ChatOpenAI
    from langgraph.prebuilt import create_react_agent

    logger.info("  [Parent Tool] delegating research for topic=%r...", topic)

    # Initialize the LLM client pointing to the Infralo gateway
    llm = ChatOpenAI(
        model=MODEL_DEPLOYMENT,
        base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
        api_key=INFRALO_API_KEY,
    )

    # Define the Child (Nested) Agent Graph
    search_tool = lc_tool(search_web)
    child_agent = create_react_agent(
        llm,
        tools=[search_tool],
    )

    # Instantiate a handler for the child agent run to capture its execution spans
    child_handler = InfraloCallbackHandler()

    # Invoke the child agent. The global HTTPX patch automatically links it under this tool!
    result = await child_agent.ainvoke(
        {"messages": [HumanMessage(content=f"Research this topic: {topic}")]},
        config={"callbacks": [child_handler]},
    )

    last_msg = result["messages"][-1].content
    logger.info("  [Parent Tool] delegation completed.")
    return last_msg


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


async def run_langgraph_complex_nesting() -> None:
    try:
        from langchain_core.messages import HumanMessage
        from langchain_core.tools import tool as lc_tool
        from langchain_openai import ChatOpenAI
        from langgraph.prebuilt import create_react_agent
    except ImportError:
        logger.error(
            "langchain-openai / langgraph not installed. "
            "Run: pip install infralo[langgraph] langchain-openai"
        )
        return

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

    # Initialize the LLM client pointing to the Infralo gateway
    llm = ChatOpenAI(
        model=MODEL_DEPLOYMENT,
        base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
        api_key=INFRALO_API_KEY,
    )

    # Define the Parent Agent Graph
    parent_tool = lc_tool(delegate_research_tool)
    parent_agent = create_react_agent(
        llm,
        tools=[parent_tool],
    )

    # Run the parent agent graph under an active Infralo trace
    with infralo.start_trace(session_id="langgraph-complex-nesting") as trace:
        print(f"Trace ID: {trace.trace_id}")

        parent_handler = InfraloCallbackHandler()

        result = await parent_agent.ainvoke(
            {"messages": [HumanMessage(content="Please research LangGraph subgraphs.")]},
            config={"callbacks": [parent_handler]},
        )

        final_response = result["messages"][-1].content
        print(f"\nFinal Parent Agent response:\n{final_response}")

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


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