#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
main.py – CLI чат‑цикл для RAG‑агента с ChromaDB и веб‑поиском через Tavily.
"""

import os
from pathlib import Path

from dotenv import load_dotenv
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.agents import AgentExecutor, create_openai_tools_agent

# --------------------------------------------------------------------------- #
# 1. Загрузка переменных окружения (TAVILY_API_KEY)
# --------------------------------------------------------------------------- #
load_dotenv()
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
if not TAVILY_API_KEY:
    raise RuntimeError("Не найдено переменной окружения TAVILY_API_KEY")

# --------------------------------------------------------------------------- #
# 2. Векторное хранилище ChromaDB
# --------------------------------------------------------------------------- #
CHROMA_DIR = Path("./chroma_db")
CHROMA_DIR.mkdir(parents=True, exist_ok=True)

embeddings = OllamaEmbeddings(model="nomic-embed-text")
vectorstore = Chroma(
    persist_directory=str(CHROMA_DIR),
    embedding_function=embeddings,
)
# --------------------------------------------------------------------------- #
# 3. Инструменты агента
# --------------------------------------------------------------------------- #

@tool("search_local_kb")
def search_local_kb(query: str, top_k: int = 5) -> str:
    """
    Семантический поиск в локальной базе знаний (ChromaDB).
    Возвращает объединённый текст найденных документов.
    """
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.invoke(query)
    if not docs:
        return "Ничего не найдено."
    # Собираем содержимое всех документов
    return "\n\n".join(doc.page_content for doc in docs)


@tool("web_search")
def web_search(query: str) -> str:
    """
    Поиск в интернете через Tavily.
    Возвращает краткое резюме результатов поиска.
    """
    from tavily import TavilyClient

    client = TavilyClient(api_key=TAVILY_API_KEY)
    results = client.search(query=query, max_results=3)
    if not results:
        return "Ничего не найдено в интернете."
    # Формируем строку с заголовками и ссылками
    snippets = []
    for r in results:
        snippets.append(f"{r['title']}\n{r['url']}")
    return "\n\n".join(snippets)


# --------------------------------------------------------------------------- #
# 4. Создание агента
# --------------------------------------------------------------------------- #

llm = ChatOllama(model="llama3", temperature=0)

tools = [search_local_kb, web_search]

agent = create_openai_tools_agent(
    llm,
    tools,
    system_message="""
Вы – интеллектуальный агент, который отвечает на вопросы пользователя.
Если вопрос относится к локальным документам (например, конспектам, заметкам),
используйте инструмент «search_local_kb».
Если требуется актуальная информация из интернета (новости, факты), используйте
инструмент «web_search».

После выполнения запроса возвращайте ответ в формате:
[Источник] <текст ответа>

Где источник – либо «chromadb», либо «tavily».
""",
)

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False)


# --------------------------------------------------------------------------- #
# 5. Чат‑цикл CLI
# --------------------------------------------------------------------------- #

def main() -> None:
    print("=== RAG‑агент с ChromaDB и Tavily ===")
    print("Введите запрос (или 'exit' для выхода):")

    while True:
        user_input = input("\n> ").strip()
        if not user_input or user_input.lower() == "exit":
            print("До свидания!")
            break

        # Запускаем агента
        try:
            result = agent_executor.invoke({"input": user_input})
        except Exception as e:
            print(f"Ошибка при обработке запроса: {e}")
            continue

        # Результат – строка вида "[Источник] ...". Если формат не соответствует,
        # просто выводим как есть.
        if isinstance(result, str) and result.startswith("["):
            try:
                source_part, answer = result.split("]", 1)
                source = source_part.strip("[")
                print(f"\nОтвет: {answer.strip()}")
                print(f"Источник: {source.lower()}")
            except ValueError:
                # Не удалось распарсить – выводим полностью
                print("\nОтвет:", result)
        else:
            print("\nОтвет:", result)


if __name__ == "__main__":
    main()