# main.py
# Полностью рабочий пример агента с RAG‑памятью на базе Qdrant и Ollama.
# Используется LangChain 1.x, create_agent, инструменты @tool, и Ollama‑LLM/embeddings.
#
# Запуск:
#   python main.py
# После запуска можно использовать команды:
#   /add <title> <content>   – добавить документ
#   /search <query> <max>    – семантический поиск
#   /quit                    – выйти
#
# Для загрузки документов из директории используйте функцию load_documents_from_dir.
#"""

import os
import sys
import textwrap
from pathlib import Path
from typing import List, Dict, Any

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.agents import create_agent, AgentExecutor, AgentToolkit, Tool
from langchain_core.messages import HumanMessage

# ---------------------------------------------------------------------------
# Конфигурация
# ---------------------------------------------------------------------------
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
QDRANT_COLLECTION = "knowledge_base"
EMBEDDING_MODEL = "nomic-embed-text"
LLM_MODEL = "llama3"

# ---------------------------------------------------------------------------
# Векторное хранилище
# ---------------------------------------------------------------------------
# Инициализируем эмбеддер и клиент Qdrant
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
vector_store = QdrantVectorStore(
    url=QDRANT_URL,
    collection_name=QDRANT_COLLECTION,
    embedding=embeddings,
)

# ---------------------------------------------------------------------------
# Чанкинг
# ---------------------------------------------------------------------------
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)

# ---------------------------------------------------------------------------
# Инструменты
# ---------------------------------------------------------------------------
@tool("search_knowledge_base", "Semantic search in the knowledge base.")
def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """Return top‑k relevant documents for a query.
    The function returns a formatted string with titles and snippets.
    """
    results = vector_store.similarity_search_with_score(query, k=max_results)
    if not results:
        return "No relevant documents found."
    formatted = []
    for doc, score in results:
        title = doc.metadata.get("title", "Untitled")
        snippet = doc.page_content[:200].replace("\n", " ")
        formatted.append(f"{title} (score: {score:.3f}): {snippet}...")
    return "\n".join(formatted)

@tool("add_to_knowledge_base", "Add a document to the knowledge base.")
def add_to_knowledge_base(content: str, title: str) -> str:
    """Chunk the content, embed, and store in Qdrant.
    Returns a confirmation message.
    """
    chunks = text_splitter.split_text(content)
    docs = []
    for i, chunk in enumerate(chunks):
        docs.append(
            {
                "page_content": chunk,
                "metadata": {"title": title, "chunk_index": i},
            }
        )
    vector_store.add_documents(docs)
    return f"Added {len(chunks)} chunks of '{title}' to the knowledge base."

# ---------------------------------------------------------------------------
# Агент
# ---------------------------------------------------------------------------
# Создаём LLM
llm = ChatOllama(model=LLM_MODEL, temperature=0.2)

# Список инструментов
tools = [search_knowledge_base, add_to_knowledge_base]

# Создаём агент
agent = create_agent(
    llm=llm,
    tools=tools,
    system_message="You are an assistant that can search and add documents to a local knowledge base. Use the provided tools.",
    verbose=True,
)

# Обёртка для выполнения
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# ---------------------------------------------------------------------------
# Загрузка документов из директории
# ---------------------------------------------------------------------------

def load_documents_from_dir(directory: str) -> None:
    """Load all .txt files from a directory into the knowledge base.
    Each file becomes a separate document with its filename as title.
    """
    path = Path(directory)
    if not path.is_dir():
        print(f"Directory {directory} does not exist.")
        return
    for file in path.glob("*.txt"):
        title = file.stem
        content = file.read_text(encoding="utf-8")
        print(f"Adding {title}...", end=" ")
        result = add_to_knowledge_base(content, title)
        print(result)

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    # Если пользователь передал путь к директории, загрузим документы
    if len(sys.argv) > 1:
        load_documents_from_dir(sys.argv[1])

    print("\n--- RAG Agent CLI ---")
    print("Commands:")
    print("  /add <title> <content>   – add a document")
    print("  /search <query> <max>    – search knowledge base")
    print("  /quit                    – exit")

    while True:
        try:
            user_input = input("\n> ")
        except (EOFError, KeyboardInterrupt):
            print("\nExiting.")
            break

        if not user_input.strip():
            continue

        if user_input.startswith("/quit"):
            print("Goodbye!")
            break

        if user_input.startswith("/add"):
            parts = user_input.split(maxsplit=2)
            if len(parts) < 3:
                print("Usage: /add <title> <content>")
                continue
            title, content = parts[1], parts[2]
            print(add_to_knowledge_base(content, title))
            continue

        if user_input.startswith("/search"):
            parts = user_input.split(maxsplit=2)
            if len(parts) < 2:
                print("Usage: /search <query> [max_results]")
                continue
            query = parts[1]
            max_results = int(parts[2]) if len(parts) > 2 else 5
            print(search_knowledge_base(query, max_results))
            continue

        # Любой другой ввод – передаём агенту
        response = agent_executor.invoke({"input": user_input})
        print(response.get("output", ""))


if __name__ == "__main__":
    main()