"""Иерархический AI-агент для планирования списка покупок (LangChain + LM Studio)."""
from __future__ import annotations

import os
import re

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

LM_STUDIO_BASE_URL = os.getenv("LM_STUDIO_BASE_URL", "http://localhost:1234/v1")
LM_STUDIO_MODEL = os.getenv("LM_STUDIO_MODEL", "local-model")


def build_llm() -> ChatOpenAI:
    return ChatOpenAI(
        model=LM_STUDIO_MODEL,
        base_url=LM_STUDIO_BASE_URL,
        api_key=SecretStr(os.getenv("OPENAI_API_KEY", "fake")),
        temperature=0.7,
    )


def _extract_table(text: str) -> str:
    """Оставляет markdown-таблицу из ответа субагента."""
    lines = [line for line in text.splitlines() if "|" in line]
    if lines:
        return "\n".join(lines)
    return text.strip()


def _build_price_subagent(llm: ChatOpenAI):
    return create_agent(
        model=llm,
        system_prompt=(
            "Ты аналитик цен на продукты питания. "
            "По названию продукта и городу оцени реалистичную цену в рублях, "
            "опираясь на типичные российские розничные цены. "
            "Ответь ТОЛЬКО одной строкой markdown-таблицы в формате:\n"
            "| Продукт | Цена (руб.) | Магазин |"
        ),
    )


def make_get_price_tool(llm: ChatOpenAI):
    price_subagent = _build_price_subagent(llm)

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

        Args:
            product: название продукта (молоко, хлеб, яблоки и т.д.)
            city: город покупателя
        """
        prompt = (
            f"Город: {city}. Продукт: {product}. "
            "Верни одну строку таблицы с реалистичной ценой и названием магазина."
        )
        result = price_subagent.invoke(
            {"messages": [{"role": "human", "content": prompt}]}
        )
        last = result["messages"][-1]
        content = getattr(last, "content", str(last))
        return _extract_table(content)

    return get_price


def format_message(message) -> str:
    content = getattr(message, "content", None)
    if content:
        return str(content)
    tool_calls = getattr(message, "tool_calls", None) or []
    if tool_calls:
        call = tool_calls[0]
        name = call.get("name") if isinstance(call, dict) else getattr(call, "name", "")
        args = call.get("args") if isinstance(call, dict) else getattr(call, "args", {})
        return f"{name}({args})"
    return str(message)


def _parse_price_row(row: str) -> int | None:
    cells = [c.strip() for c in row.strip("|").split("|") if c.strip()]
    if len(cells) < 2:
        return None
    price_raw = cells[1]
    numbers = re.findall(r"\d+", price_raw.replace(" ", ""))
    if not numbers:
        return None
    return int(numbers[0])


def run_shopping_assistant() -> None:
    llm = build_llm()
    get_price = make_get_price_tool(llm)

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

    question = (
        "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    )
    answer = main_agent.invoke(
        {"messages": [{"role": "human", "content": question}]}
    )

    print("--- Цепочка сообщений ---")
    for message in answer["messages"]:
        formatted = format_message(message)
        if formatted.strip():
            print("---")
            print(formatted)

    print("\n--- Итог ---")
    final = answer["messages"][-1]
    print(getattr(final, "content", final))


if __name__ == "__main__":
    run_shopping_assistant()