from __future__ import annotations

import argparse
import os
from typing import Any

from langchain_core.messages import AIMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent


@tool
def get_demo_price(product: str, city: str = "Казань") -> str:
    """Return a deterministic demo price for the requested product."""
    prices = {
        "молоко": "89 рублей",
        "хлеб": "54 рубля",
        "сыр": "219 рублей",
    }
    return f"{product} в городе {city}: {prices.get(product.lower(), 'цена не найдена')}"


def build_agent():
    llm = ChatOpenAI(
        model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"),
        base_url=os.getenv("OPENAI_BASE_URL") or None,
        api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
        temperature=0,
        streaming=True,
    )
    return create_react_agent(llm, tools=[get_demo_price])


def format_message(message: AIMessage) -> str:
    if message.content:
        return str(message.content)
    if message.tool_calls:
        call = message.tool_calls[0]
        return f"{call['name']}({call['args']})"
    return ""


def format_chunk_message(chunk: tuple[Any, dict[str, Any]], state: dict[str, int]) -> None:
    message, meta = chunk
    current_step = int(meta.get("langgraph_step", state["step"]))
    if current_step != state["step"]:
        state["step"] = current_step
        print("\n --- --- --- \n")
    if getattr(message, "content", None):
        print(message.content, end="", flush=True)


def run_stream(question: str) -> None:
    agent = build_agent()
    stream = agent.stream(
        {"messages": [{"role": "human", "content": question}]},
        stream_mode=["messages", "updates"],
    )
    state = {"step": 1}
    for chunk_type, chunk_data in stream:
        if chunk_type == "messages":
            format_chunk_message(chunk_data, state)
        elif chunk_type == "updates" and chunk_data.get("model"):
            last_message = chunk_data["model"]["messages"][-1]
            formatted = format_message(last_message)
            if formatted:
                print(f"\n{formatted}")

def main() -> None:
    parser = argparse.ArgumentParser(description="Streaming AI-agent demo")
    parser.add_argument("question", nargs="*", help="User question")
    args = parser.parse_args()

    question = " ".join(args.question).strip() or "Сколько стоит молоко и хлеб в Казани?"
    run_stream(question)


if __name__ == "__main__":
    main()