import os
from typing import Dict, Any

from langchain_ollama import ChatOllama
from nomic_embed_text import NomicEmbedText
from langchain.tools import BaseTool
from langchain.schema import HumanMessage, AIMessage, SystemMessage
from langchain.agents import AgentExecutor, initialize_agent, load_tools
from langchain.prompts import ChatPromptTemplate

# ---------- Настройки 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]

# ---------- Создание агента ----------
prompt_template = ChatPromptTemplate.from_messages(
    [
        SystemMessage(content="Ты помощник. Используй инструменты при необходимости."),
        ("human", "{input}"),
    ]
)

agent_executor = initialize_agent(
    tools,
    llm,
    agent="openai-tools",
    verbose=False,
    prompt=prompt_template,
)

# ---------- Функции форматирования ----------
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 = {"input": user_input}
    # В LangChain 0.2 AgentExecutor имеет метод stream()
    stream = agent_executor.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)