"""Интерактивный RAG-агент (ChromaDB + Ollama + LangChain)."""
from __future__ import annotations

import os

from langchain.agents import create_agent
from langchain_ollama import ChatOllama

from rag_tools import add_to_knowledge_base, search_knowledge_base

LLM_MODEL = os.getenv("OLLAMA_LLM_MODEL", "llama3")
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")

SYSTEM_PROMPT = """\
Ты ассистент с доступом к базе знаний.
Перед ответом на фактические вопросы используй search_knowledge_base.
Если пользователь просит сохранить информацию — вызывай add_to_knowledge_base.
Отвечай на русском, кратко и по существу.
"""


def build_agent():
    llm = ChatOllama(
        model=LLM_MODEL,
        base_url=OLLAMA_BASE_URL,
        temperature=0.2,
    )
    return create_agent(
        model=llm,
        tools=[search_knowledge_base, add_to_knowledge_base],
        system_prompt=SYSTEM_PROMPT,
    )


def run_interactive_cli() -> None:
    agent = build_agent()
    print("RAG-агент. Команды: /add, /search, /quit")
    print("  /add <заголовок> | <текст>")
    print("  /search <запрос>")

    while True:
        user_input = input("\n> ").strip()
        if not user_input:
            continue
        if user_input == "/quit":
            print("Выход.")
            break

        if user_input.startswith("/search "):
            query = user_input[len("/search ") :].strip()
            print(search_knowledge_base.invoke({"query": query, "max_results": 5}))
            continue

        if user_input.startswith("/add "):
            payload = user_input[len("/add ") :].strip()
            if "|" not in payload:
                print("Формат: /add <заголовок> | <текст>")
                continue
            title, content = (part.strip() for part in payload.split("|", 1))
            print(add_to_knowledge_base.invoke({"content": content, "title": title}))
            continue

        result = agent.invoke(
            {"messages": [{"role": "human", "content": user_input}]}
        )
        answer = result["messages"][-1]
        print(getattr(answer, "content", answer))


if __name__ == "__main__":
    run_interactive_cli()