import os
from typing import Dict, Any

from langgraph.graph import StateGraph, START
from langgraph.prebuilt.tool_executor import ToolExecutorNode
from langgraph.checkpoint.memory import MemorySaver
from langchain_ollama import ChatOllama
from nomic_embed_text import NomicEmbedText
from langchain.tools import BaseTool
from langchain.schema import HumanMessage, AIMessage, SystemMessage

# ---------- Настройки LLM и эмбеддингов ----------
LLM_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
EMBEDDING_MODEL = os.getenv("NOMIC_EMBEDDING_MODEL", "nomic-embed-text")

llm = ChatOllama(model=LLM_MODEL, temperature=0.2)
embedding = NomicEmbedText()

# ---------- Пример инструмента ----------
class DummyTool(BaseTool):
    name: str = "get_price"
    description: str = (
        "Получает цену товара в указанном городе. "
        "Аргументы: {'product': 'название', 'city': 'город'}."
    )

    def _run(self, product: str, city: str) -> str:
        # В реальном коде здесь будет запрос к API
        return f"Цена {product} в {city}: 100₽"

tool = DummyTool()
tools = [tool]

# ---------- Создание агента ----------
def create_agent() -> StateGraph:
    graph = StateGraph()

    def agent(state: Dict[str, Any]) -> Dict[str, Any]:
        messages = state.get("messages", [])
        # Добавляем системное сообщение с инструкцией
        system_msg = SystemMessage(
            content="Ты помощник. Используй инструменты при необходимости."
        )
        all_messages = [system_msg] + messages

        # Запускаем LLM в режиме stream
        response = llm.invoke(all_messages, stream=True)
        return {"messages": [response]}

    graph.add_node("agent", agent)

    def tool_executor(state: Dict[str, Any]) -> Dict[str, Any]:
        last_msg = state["messages"][-1]
        if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
            # Выполняем первый вызов инструмента
            call = last_msg.tool_calls[0]
            name = call["name"]
            args = eval(call["args"])
            result = tools_dict[name].invoke(args)
            # Создаём сообщение с результатом
            tool_msg = AIMessage(
                content=result,
                tool_calls=[call],
            )
            return {"messages": [tool_msg]}
        return {}

    graph.add_node("tool_executor", ToolExecutorNode(tools))
    graph.set_entry_point("agent")
    graph.add_edge(START, "agent")
    graph.add_conditional_edges(
        "agent",
        lambda x: "messages" in x and hasattr(x["messages"][-1], "tool_calls"),
        {
            True: "tool_executor",
            False: START,
        },
    )
    graph.set_finish_node("agent")

    return graph.compile(checkpointer=MemorySaver())

tools_dict = {t.name: t for t in tools}
graph = create_agent()

# ---------- Функции форматирования ----------
def format_message(message) -> str:
    if message.content:
        return message.content
    # Если нет content, выводим вызов инструмента
    call = message.tool_calls[0]
    return f"{call['name']}({call['args']})"

step = 1

def format_chunk_message(chunk):
    global step
    message, meta = chunk
    if meta.get("langgraph_step") != step:
        step = meta["langgraph_step"]
        print("\n --- --- --- \n")
    if message.content:
        print(message.content, end="", flush=True)

# ---------- Запуск агента в режиме stream ----------
def run_agent(user_input: str):
    global step
    step = 1
    # Инициализируем состояние с пользовательским сообщением
    init_state = {"messages": [HumanMessage(content=user_input)]}
    stream = graph.stream(init_state, stream_mode=["messages", "updates"])

    for chunk_type, chunk_data in stream:
        if chunk_type == "messages":
            format_chunk_message(chunk_data)
        elif chunk_type == "updates":
            # При завершении шага выводим итоговое сообщение
            model_info = chunk_data.get("model")
            if model_info:
                last_msg = model_info["messages"][-1]
                print("\n" + format_message(last_msg) + "\n")

if __name__ == "__main__":
    user_query = input("Введите запрос: ")
    run_agent(user_query)