from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from pydantic import SecretStr


llm = ChatOpenAI(
    model="local-model",
    base_url="http://localhost:1234/v1",
    api_key=SecretStr("fake"),
    temperature=0.7,
)


@tool
def get_price(product: str, city: str) -> str:
    """Возвращает примерную цену продукта в указанном городе."""

    price_agent = create_agent(
        model=llm,
        tools=[],
        system_prompt=(
            "Ты субагент оценки цен. Верни одну markdown-таблицу "
            "с колонками: Продукт, Цена (руб.), Магазин."
        ),
    )
    answer = price_agent.invoke(
        {
            "messages": [
                {
                    "role": "human",
                    "content": (
                        f"Оцени реалистичную цену продукта '{product}' "
                        f"в городе {city}, опираясь на исторические данные о ценах."
                    ),
                }
            ]
        }
    )
    return answer["messages"][-1].content


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


def stream_shopping_agent() -> None:
    memory = MemorySaver()
    agent = create_agent(
        model=llm,
        tools=[get_price],
        system_prompt="Ты помощник по планированию покупок",
        checkpointer=memory,
    )

    stream = agent.stream(
        {
            "messages": [
                {
                    "role": "human",
                    "content": (
                        "Помоги составить список покупок: молоко, хлеб, яблоки. "
                        "Я нахожусь в Казани."
                    ),
                }
            ]
        },
        stream_mode=["messages", "updates"],
        config={"configurable": {"thread_id": "shopping-stream-demo"}},
    )

    step = None
    for chunk_type, chunk_data in stream:
        if chunk_type == "messages":
            message, meta = chunk_data
            current_step = meta.get("langgraph_step")
            if current_step != step:
                step = current_step
                print("\n --- --- --- \n")
            if message.content:
                print(message.content, end="", flush=True)

        if chunk_type == "updates" and chunk_data.get("model"):
            last_message = chunk_data["model"]["messages"][-1]
            formatted = format_message(last_message)
            if formatted:
                print(formatted)


if __name__ == "__main__":
    stream_shopping_agent()