import os
from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.agents import create_agent, AgentExecutor, Tool
from langchain.schema import Document

# --------------------------------------------------------------------------- #
# Конфигурация подключения к Qdrant и Ollama
# --------------------------------------------------------------------------- #

QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")

EMBEDDING_MODEL = "nomic-embed-text"
LLM_MODEL = "llama3"

# --------------------------------------------------------------------------- #
# Векторное хранилище
# --------------------------------------------------------------------------- #

class RagVectorStore:
    """Обёртка над QdrantVectorStore с автоматическим эмбеддингом Ollama."""

    def __init__(self, collection_name: str = COLLECTION_NAME):
        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
        self.store = QdrantVectorStore(
            url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
            collection_name=collection_name,
            embeddings=self.embeddings,
        )
        # Создаём коллекцию, если её ещё нет
        if not self.store.client.get_collection(collection_name).exists:
            self.store.create_collection()

    def add_documents(self, docs: List[Document]) -> None:
        """Добавляет список документов в хранилище."""
        self.store.add_documents(docs)

    def search(
        self,
        query: str,
        k: int = 5,
    ) -> List[tuple[str, float]]:
        """Семантический поиск по запросу. Возвращает кортежи (текст, релевантность)."""
        results = self.store.similarity_search_with_score(query, k=k)
        return [(doc.page_content, score) for doc, score in results]

# --------------------------------------------------------------------------- #
# Чанкинг документов
# --------------------------------------------------------------------------- #

def chunk_document(text: str, title: str) -> List[Document]:
    """Разбивает текст на чанки и возвращает список Document."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
    )
    chunks = splitter.split_text(text)
    return [
        Document(page_content=chunk, metadata={"title": title})
        for chunk in chunks
    ]

# --------------------------------------------------------------------------- #
# Инструменты для агента
# --------------------------------------------------------------------------- #

@tool("search_knowledge_base")
def search_knowledge_base(
    query: str,
    max_results: int = 5,
) -> List[str]:
    """
    Семантический поиск в базе знаний.

    Parameters:
        query (str): поисковый запрос.
        max_results (int): максимальное количество результатов.

    Returns:
        list of strings: найденные тексты с оценкой релевантности.
    """
    results = rag_store.search(query, k=max_results)
    return [f"{text} (score={score:.4f})" for text, score in results]

@tool("add_to_knowledge_base")
def add_to_knowledge_base(
    content: str,
    title: str,
) -> str:
    """
    Добавляет новый документ в базу знаний.

    Parameters:
        content (str): полный текст документа.
        title (str): заголовок/название документа.

    Returns:
        str: сообщение об успехе.
    """
    docs = chunk_document(content, title)
    rag_store.add_documents(docs)
    return f"Документ '{title}' успешно добавлен в базу знаний."

# --------------------------------------------------------------------------- #
# Инициализация хранилища
# --------------------------------------------------------------------------- #

rag_store = RagVectorStore()

# --------------------------------------------------------------------------- #
# Создание агента с RAG‑интеграцией
# --------------------------------------------------------------------------- #

llm = OllamaLLM(model=LLM_MODEL)

tools: List[Tool] = [
    Tool.from_function(
        func=search_knowledge_base,
        name="search_knowledge_base",
        description="Используйте для поиска информации в базе знаний.",
    ),
    Tool.from_function(
        func=add_to_knowledge_base,
        name="add_to_knowledge_base",
        description="Добавьте новый документ в базу знаний.",
    ),
]

agent = create_agent(
    llm=llm,
    tools=tools,
    system_message=(
        "Вы — интеллектуальный агент, способный искать и добавлять информацию "
        "в локальную базу знаний. При необходимости используйте инструменты "
        "'search_knowledge_base' и 'add_to_knowledge_base'."
    ),
)

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

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

def load_documents_from_directory(directory: str) -> None:
    """Загружает все текстовые файлы из указанной папки в базу знаний."""
    path = Path(directory)
    for file_path in path.rglob("*"):
        if file_path.suffix.lower() in {".txt", ".md"}:
            text = file_path.read_text(encoding="utf-8")
            title = file_path.stem
            add_to_knowledge_base(content=text, title=title)

# --------------------------------------------------------------------------- #
# Интерактивный клиент
# --------------------------------------------------------------------------- #

def interactive_cli():
    """Простой CLI для взаимодействия с агентом."""
    print("=== RAG Агент ===")
    print("Команды: /add <path>, /search <query>, /quit")
    while True:
        try:
            user_input = input("> ").strip()
        except (EOFError, KeyboardInterrupt):
            break

        if not user_input:
            continue

        if user_input.lower() == "/quit":
            print("Выход.")
            break
        elif user_input.startswith("/add"):
            _, *path_parts = user_input.split(maxsplit=1)
            if path_parts:
                file_path = Path(path_parts[0])
                if file_path.is_file():
                    content = file_path.read_text(encoding="utf-8")
                    title = file_path.stem
                    print(add_to_knowledge_base(content, title))
                else:
                    print(f"Файл не найден: {file_path}")
            else:
                print("Укажите путь к файлу.")
        elif user_input.startswith("/search"):
            _, *query_parts = user_input.split(maxsplit=1)
            if query_parts:
                query = query_parts[0]
                results = search_knowledge_base(query, max_results=3)
                for r in results:
                    print(r)
            else:
                print("Укажите поисковый запрос.")
        else:
            # Любой другой ввод передаём агенту
            response = agent_executor.invoke({"input": user_input})
            print(response["output"])

if __name__ == "__main__":
    interactive_cli()