import os
import asyncio
from langchain_openai import ChatOpenAI
from langchain.tools import tool
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage

# --- LLM setup ------------------------------------------------------------
# Connect to local LM Studio server (OpenAI compatible API)
llm = ChatOpenAI(
    model="<название модели в LM Studio>",
    base_url="http://localhost:1234/v1",
    api_key=os.getenv("OPENAI_API_KEY", "fake"),
    temperature=0.7,
)

# --- Sub‑agent for price generation -------------------------------------
# This sub‑agent is created inside the tool and is responsible for
# producing a realistic price table for a single product.
sub_agent = create_agent(
    model=llm,
    tools=[],  # no external tools needed for the sub‑agent
    system_prompt="Ты генератор цены. На основе исторических данных выдавай таблицу с ценой и магазином.",
)

# --- Tool definition -----------------------------------------------------
@tool
def get_price(product: str, city: str) -> str:
    """Получить цену продукта в указанном городе.

    Возвращает строку в формате Markdown‑таблицы:
    | Продукт | Цена (руб.) | Магазин |
    """
    # Формируем запрос для суб‑агента
    prompt = f"\n\nНайди примерную цену на {product} в городе {city}.\n\nВыведи результат в виде таблицы:\n| Продукт | Цена (руб.) | Магазин |"
    # Запускаем суб‑агента
    result = sub_agent.invoke({"messages": [HumanMessage(content=prompt)]})
    # sub_agent возвращает dict с ключом 'messages'
    # Берём последний текстовый ответ
    last_msg = result["messages"][-1]
    return last_msg.content.strip()

# --- Main agent ----------------------------------------------------------
agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты помощник по планированию покупок.",
)

# --- Helper for pretty printing ----------------------------------------

def format_message(message):
    if hasattr(message, "content") and message.content:
        return message.content
    if hasattr(message, "tool_calls") and message.tool_calls:
        call = message.tool_calls[0]
        return f"{call['name']}({call['args']})"
    return ""

# --- Main execution ------------------------------------------------------
async def main():
    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=user_query)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    # Выводим все сообщения
    for msg in result["messages"]:
        print(format_message(msg))

if __name__ == "__main__":
    asyncio.run(main())