"""RAG-агент с ChromaDB (Ollama) и веб-поиском (Tavily). CLI."""
from __future__ import annotations

from pathlib import Path

from dotenv import load_dotenv
from langchain_ollama import ChatOllama
from langchain.agents import AgentExecutor, create_react_agent
from langchain import hub

from vectorstore import create_vectorstore, load_documents
from tools import init_tools, search_local_kb, web_search

load_dotenv()

# ── LLM: Ollama llama3 ────────────────────────────────────────────────────
llm = ChatOllama(model="llama3", temperature=0)

# ── Векторное хранилище ───────────────────────────────────────────────────
CHROMA_DIR = "./chroma_db"
vectorstore = create_vectorstore(CHROMA_DIR)
init_tools(vectorstore)

# Автозагрузка документов из documents/ (если есть и БД пустая)
docs_dir = Path("documents")
if docs_dir.exists() and not any(Path(CHROMA_DIR).rglob("*.sqlite3")):
    count = load_documents(str(docs_dir), vectorstore)
    print(f"[init] Загружено {count} чанков из {docs_dir}/")

# ── ReAct-агент ───────────────────────────────────────────────────────────
tools = [search_local_kb, web_search]
prompt = hub.pull("hwchase17/react")
agent = create_react_agent(llm, tools, prompt)
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    max_iterations=10,
    handle_parsing_errors=True,
)

ROUTING_HINT = (
    "Для вопросов по локальным документам/конспектам используй search_local_kb. "
    "Для актуальных новостей и фактов из интернета — web_search. "
    "В ответе всегда укажи источник: chromadb или tavily."
)


def ask(question: str) -> str:
    """Задать вопрос агенту."""
    result = agent_executor.invoke({"input": f"{ROUTING_HINT}

Вопрос: {question}"})
    return result.get("output", str(result))


if __name__ == "__main__":
    print("RAG-агент: Ollama llama3 + nomic-embed-text + ChromaDB + Tavily")
    print("Введите 'exit' для выхода.")
    print("-" * 60)
    while True:
        try:
            user_input = input("
Запрос: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("
Завершение.")
            break
        if user_input.lower() in ("exit", "quit", "выход"):
            print("До свидания!")
            break
        if not user_input:
            continue
        answer = ask(user_input)
        print(f"
{answer}")