"""
Main entry point for the Stream‑mode LangChain agent.

The script demonstrates how to replace a single ``invoke`` call with a streaming
``stream`` interface.  The output is printed token by token so that the user can
see the agent “think” in real time.

Three example invocations are provided:
1. A simple question that requires no tool calls.
2. A request that triggers the ``get_price`` tool.
3. A multi‑step conversation that uses the same tool twice.
"""

from __future__ import annotations

import os
from typing import Any, Dict, Tuple

# ---------------------------------------------------------------------------
# LLM configuration – BroJS provider
# ---------------------------------------------------------------------------
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
    api_key=os.getenv("JOURNAL_MCP_PAT"),
    temperature=0.0,
)

# ---------------------------------------------------------------------------
# Tool definition – a tiny mock price lookup
# ---------------------------------------------------------------------------
from langchain.tools import tool

@tool
def get_price(product: str, city: str) -> str:
    """Return a fake price for *product* in *city*.

    The function is intentionally simple; it only demonstrates how the agent
    can call tools while streaming.  In a real project this would query an API
    or database.
    """
    prices = {
        ("milk", "kazan"): "89",
        ("bread", "kazan"): "45",
        ("coffee", "moscow"): "120",
    }
    key = (product.lower(), city.lower())
    price = prices.get(key, "unknown")
    return f"{price} rubles"

# ---------------------------------------------------------------------------
# Agent construction – zero‑shot react description
# ---------------------------------------------------------------------------
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage

agent = create_agent(
    llm=llm,
    tools=[get_price],
    system_prompt="You are a helpful assistant that can call the get_price tool.",
)

# ---------------------------------------------------------------------------
# Streaming helper functions
# ---------------------------------------------------------------------------
def format_message(message: Any) -> str:
    """Return human‑readable representation of a message.

    If the message contains content we return it directly.  Otherwise we
    construct a string that shows the tool call.
    """
    if getattr(message, "content", None):
        return message.content
    # Tool call – ``message.tool_calls`` is a list of dicts
    calls = [f"{c['name']}({c['args']})" for c in message.tool_calls]
    return ", ".join(calls)

# ---------------------------------------------------------------------------
# Core streaming logic
# ---------------------------------------------------------------------------
def stream_and_print(prompt: str) -> None:
    """Invoke the agent with ``stream`` and print tokens as they arrive.

    Parameters
    ----------
    prompt: str
        The user message to send to the agent.
    """
    # Start streaming – we want both token‑level messages and state updates
    stream = agent.stream(
        {"messages": [HumanMessage(content=prompt)]},
        stream_mode=["messages", "updates"],
    )

    current_step: int | None = None

    for chunk_type, chunk_data in stream:
        if chunk_type == "messages":
            # ``chunk_data`` is a tuple (message, meta)
            message, meta = chunk_data  # type: ignore[assignment]
            step = meta.get("langgraph_step")
            if current_step != step:
                current_step = step
                print("\n--- Step {} ---\n".format(step), end="", flush=True)
            if message.content:
                print(message.content, end="", flush=True)
        elif chunk_type == "updates":
            # ``chunk_data`` contains the finished state of a step
            model = chunk_data.get("model")
            if model and model["messages"]:
                last_msg = model["messages"][-1]
                print(format_message(last_msg), end="", flush=True)
    print("\n--- End ---\n")

# ---------------------------------------------------------------------------
# Example usage – three distinct scenarios
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    examples = [
        "What is the capital of France?",
        "How much does milk cost in Kazan?",
        "I need prices for bread and coffee in Kazan and Moscow.",
    ]

    for i, ex in enumerate(examples, 1):
        print(f"\n=== Example {i} ===")
        stream_and_print(ex)

# End of file