"""Агент с памятью разговора и подтверждением вызовов инструментов."""
from __future__ import annotations

import os

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
from rich.console import Console

console = Console()

LM_STUDIO_BASE_URL = os.getenv("LM_STUDIO_BASE_URL", "http://localhost:1234/v1")
LM_STUDIO_MODEL = os.getenv("LM_STUDIO_MODEL", "local-model")
STEP_SEPARATOR = "\n --- --- --- \n"
THREAD_ID = "разговор-1"


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:
    lines = [line for line in text.splitlines() if "|" in line]
    return "\n".join(lines) if lines else text.strip()


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


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

    @tool
    def get_price(product: str, city: str, date: str = "") -> str:
        """Возвращает примерную цену продукта в городе (опционально с учётом даты/периода).

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

    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 _tool_call_repr(state) -> str:
    messages = state.values.get("messages", [])
    if not messages:
        return "unknown_tool()"
    last = messages[-1]
    tool_calls = getattr(last, "tool_calls", None) or []
    if not tool_calls:
        return "unknown_tool()"
    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})'


def build_agent(llm: ChatOpenAI):
    memory = MemorySaver()
    get_price = make_get_price_tool(llm)
    return create_agent(
        model=llm,
        tools=[get_price],
        system_prompt="Ты помощник по планированию покупок",
        checkpointer=memory,
        interrupt_before=["tools"],
    )


def run_chat() -> None:
    llm = build_llm()
    agent = build_agent(llm)
    config = {"configurable": {"thread_id": THREAD_ID}}
    step_holder = {"step": 1}

    def format_chunk_message(chunk) -> None:
        message, meta = chunk
        graph_step = meta.get("langgraph_step", step_holder["step"])
        if graph_step != step_holder["step"]:
            step_holder["step"] = graph_step
            console.print(STEP_SEPARATOR, end="")
        if message.content:
            console.print(message.content, end="")

    def ask_and_run(user_input, cfg) -> None:
        stream = agent.stream(
            user_input,
            config=cfg,
            stream_mode=["messages", "updates"],
        )

        for chunk in stream:
            state = agent.get_state(cfg)
            chunk_type, chunk_data = chunk

            if chunk_type == "messages":
                format_chunk_message(chunk_data)

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

                if "__interrupt__" in chunk_data and state.next == ("tools",):
                    tool_repr = _tool_call_repr(state)
                    console.print(STEP_SEPARATOR, end="")
                    console.print(tool_repr)
                    console.print(
                        f"Агент хочет вызвать утилиту {tool_repr}"
                    )
                    answer = input("Разрешить? (Y/n): ")
                    if answer.lower().strip() == "y":
                        ask_and_run(None, cfg)
                    else:
                        console.print("Отменено")
                    return

        console.print()

    console.print(
        "Чат с агентом (память + подтверждение инструментов). "
        "Введите exit для выхода."
    )

    while True:
        user_input = input("\nВы: ")
        if user_input.strip().lower() == "exit":
            break
        ask_and_run(
            {"messages": [{"role": "human", "content": user_input}]},
            config,
        )


if __name__ == "__main__":
    run_chat()