"""
LangGraph / LangChain Integration Example — infralo-sdk

Demonstrates:
  1. Using InfraloCallbackHandler with a LangGraph ReAct agent.
  2. Using InfraloCallbackHandler with a plain LangChain chain.
  3. Tool span capture with inputs, outputs, and error handling.

Requirements:
    pip install infralo[langgraph] langchain-openai

Run:
    cd examples/langgraph
    uv run langgraph_integration.py
    # or: python langgraph_integration.py
"""

import asyncio
import logging
import os

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

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

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 ─────────────────────────────────────────────────────────────


def search_web(query: str) -> str:
    """Search the web and return relevant results."""
    logger.info("  [tool] search_web(query=%r)", query)
    return f"Web results for '{query}': item 1, item 2, item 3."


def calculate(expression: str) -> str:
    """Evaluate a simple mathematical expression safely."""
    logger.info("  [tool] calculate(expression=%r)", expression)
    try:
        result = eval(expression, {"__builtins__": {}})  # noqa: S307, PGH001
        return str(result)
    except Exception as exc:
        raise ValueError(f"Cannot evaluate: {exc}") from exc


# ── 3. LangGraph ReAct Agent run ─────────────────────────────────────────────


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

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

    # Wrap tools with LangChain's @tool decorator so LangGraph discovers them
    search_tool = lc_tool(search_web)
    calc_tool = lc_tool(calculate)

    llm = ChatOpenAI(
        model=MODEL_DEPLOYMENT,
        base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
        api_key=INFRALO_API_KEY,
    )

    agent = create_react_agent(llm, [search_tool, calc_tool])

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

        # A new handler per invocation is the safest pattern
        handler = InfraloCallbackHandler()

        result = await agent.ainvoke(
            {"messages": [HumanMessage(content="What is 42 * 7? Also search for LangGraph news.")]},
            config={"callbacks": [handler]},
        )

        last_message = result["messages"][-1]
        print(f"\nAgent response:\n{last_message.content}")

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


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