import random
from typing import Literal

from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from pydantic import SecretStr

# ---------------------------------------------------------------------------
# 1. Подключение к локальной модели (LM Studio)
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
    model="<название_модели>",  # замените на имя модели, запущенной в LM Studio
    base_url="http://localhost:1234/v1",
    api_key=SecretStr("fake"),
    temperature=0.7,
)

# ---------------------------------------------------------------------------
# 2. Инструмент с суб‑агентом: get_price
# ---------------------------------------------------------------------------
@tool
def get_price(product: str, city: str) -> str:
    """Возвращает таблицу с примерной ценой продукта в указанном городе.

    Суб‑агент генерирует цену и форматирует её в виде Markdown‑таблицы.
    """

    # ------------------- суб‑агент -------------------
    # Простой суб‑агент, использующий тот же LLM, но с более специфическим
    # промптом, чтобы сгенерировать реалистичную цену.
    sub_prompt = (
        f"Ты экономический аналитик. Дай примерную цену продукта '{product}' в городе '{city}'. "
        "Ответ дай в виде строки цены (число) и названия магазина, без лишних символов. "
        "Если нет данных – придумай разумную цену в диапазоне 30‑200 рублей."
    )

    sub_agent = create_agent(
        model=llm,
        tools=[],  # суб‑агент не использует внешних инструментов
        system_prompt=sub_prompt,
    )

    # Запрос к суб‑агенту
    sub_answer = sub_agent.invoke({"messages": [{"role": "human", "content": "Цена?"}]})
    # Получаем текст ответа
    price_text = sub_answer["messages"][-1]["content"].strip()

    # Пытаемся извлечь число из ответа; если не получится – генерируем случайно
    try:
        # ищем первое число в строке
        import re
        match = re.search(r"\d+", price_text)
        price = int(match.group()) if match else random.randint(30, 200)
    except Exception:
        price = random.randint(30, 200)

    # Выбираем магазин произвольно
    stores = ["Магнит", "Пятёрочка", "Перекрёсток", "Ашан", "Дикси"]
    store = random.choice(stores)

    # Формируем таблицу Markdown
    table = (
        "| Продукт | Цена (руб.) | Магазин |\n"
        "|---------|-------------|---------|\n"
        f"| {product.capitalize()} | {price} | {store} |"
    )
    return table

# ---------------------------------------------------------------------------
# 3. Главный агент
# ---------------------------------------------------------------------------
main_agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price, чтобы узнать цены, и подсчитай итоговую стоимость.",
)

# ---------------------------------------------------------------------------
# 4. Утилита для подсчёта итоговой стоимости из полученных таблиц
# ---------------------------------------------------------------------------
def parse_price(table: str) -> int:
    """Извлекает цену из строки таблицы Markdown."""
    import re
    match = re.search(r"\|\s*[^|]+\|\s*(\d+)\s*\|", table)
    return int(match.group(1)) if match else 0

# ---------------------------------------------------------------------------
# 5. Основная функция: принимает запрос пользователя и выводит цепочку сообщений
# ---------------------------------------------------------------------------
def run_shopping_assistant(user_query: str):
    """Запускает главного агента с пользовательским запросом и выводит все шаги."""
    # Запускаем агент в режиме stream, чтобы увидеть промежуточные вызовы
    stream = main_agent.stream({"messages": [{"role": "human", "content": user_query}]},
                              stream_mode=["messages", "updates"])
    step = 0
    for chunk_type, chunk_data in stream:
        if chunk_type == "messages":
            message, meta = chunk_data
            if meta.get("langgraph_step") != step:
                step = meta.get("langgraph_step")
                print("\n--- step", step, "---\n")
            if message.get("content"):
                print(message["content"], end="")
        elif chunk_type == "updates":
            # При завершении шага модели мы получаем полное сообщение
            if "model" in chunk_data:
                last_msg = chunk_data["model"]["messages"][-1]
                if last_msg.get("content"):
                    print(last_msg["content"], end="")
                elif last_msg.get("tool_calls"):
                    # Выводим вызов инструмента
                    tc = last_msg["tool_calls"][0]
                    print(f"{tc['name']}({tc['args']})")
    print("\n\n--- Финальный вывод ---\n")
    # После завершения получаем полную историю
    final = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
    # Вычислим итоговую стоимость, если в ответе есть таблицы
    total = 0
    for msg in final["messages"]:
        if msg.get("content") and "|" in msg["content"]:
            total += parse_price(msg["content"])
    if total:
        print(f"Итоговая стоимость: {total} руб.")

# ---------------------------------------------------------------------------
# 6. Тестовый запуск (можно закомментировать при импортировании)
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    run_shopping_assistant(query)