"""A real LangChain agent with Toolcue: three turns with a real model, showing which tools the model got each time.

    pip install -r requirements.txt          # toolcue[langchain] and a model provider
    export TOOLCUE_URL=... TOOLCUE_TOKEN=... TOOLCUE_ADMIN_TOKEN=...
    export ANTHROPIC_API_KEY=...                         # or OPENAI_API_KEY=... and OPENAI_MODEL=...
    python agent.py
"""

import os
import sys
from dataclasses import dataclass

from langchain.agents import create_agent
from langchain.agents.middleware import AgentMiddleware
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from toolcue.langchain import ToolcueMiddleware, publish_tools

from billing_tools import EXTRAS, TOOLS

AGENT = "acme-billing"


@dataclass
class Session:
    roles: list[str]


class ShowTools(AgentMiddleware):
    """Prints the tools each model call receives (after Toolcue filtered them)."""

    def wrap_model_call(self, request, handler):
        print(f"  model gets {len(request.tools)}/{len(TOOLS) + 1} tools:",
              ", ".join(sorted(t.name for t in request.tools if hasattr(t, "name"))))
        return handler(request)


def model_name() -> str:
    if os.environ.get("ANTHROPIC_API_KEY"):
        return "anthropic:claude-opus-5-5"
    if os.environ.get("OPENAI_API_KEY") and os.environ.get("OPENAI_MODEL"):
        return "openai:" + os.environ["OPENAI_MODEL"]
    sys.exit("Set ANTHROPIC_API_KEY, or OPENAI_API_KEY and OPENAI_MODEL.")


def main():
    summary = publish_tools(AGENT, TOOLS, EXTRAS)
    print(f"published {summary['tools']} tools, {summary['edges']} dependency links")

    agent = create_agent(
        model_name(),
        tools=TOOLS,
        system_prompt="You are the billing assistant of Acme. Use the tools; never invent ids.",
        middleware=[ToolcueMiddleware(AGENT, user=lambda session: {"roles": session.roles}), ShowTools()],
        context_schema=Session,
        checkpointer=InMemorySaver(),
    )
    thread = {"configurable": {"thread_id": "demo"}}
    for text in ["Refund the last invoice for Globex", "thanks! what plan are they on now?",
                 "ok, move them to the scale plan"]:
        print(f"\n> {text}")
        result = agent.invoke({"messages": [HumanMessage(text)]}, thread, context=Session(roles=["agent"]))
        calls = [c["name"] for m in result["messages"] for c in getattr(m, "tool_calls", [])]
        print("  tool calls so far:", ", ".join(calls) or "none")
        print("  answer:", result["messages"][-1].text)


if __name__ == "__main__":
    main()