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

"""
solution.py – основной скрипт с реализацией RAG‑агента и инструментов.
"""

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

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

# --------------------------------------------------------------------------- #
# Конфигурация
# --------------------------------------------------------------------------- #

QDRANT_HOST = "localhost"
QDRANT_PORT = 6333
COLLECTION_NAME = "knowledge_base"

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

MAX_CHUNK_SIZE = 1000   # символов
CHUNK_OVERLAP = 200     # символов


# --------------------------------------------------------------------------- #
# Векторное хранилище (обёртка над QdrantVectorStore)
# --------------------------------------------------------------------------- #

class KnowledgeBase:
    """
    Класс, инкапсулирующий работу с Qdrant и Ollama.
    """

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

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

    def similarity_search(
        self,
        query: str,
        k: int = 5
    ) -> List[Dict]:
        """
        Семантический поиск по запросу. Возвращает списки словарей с полями:
        'content', 'metadata'.
        """
        results = self.store.similarity_search(query, k=k)
        return [
            {"content": doc.page_content,
             "metadata": doc.metadata}
            for doc in results
        ]


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

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


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

kb = KnowledgeBase()


@tool("search_knowledge_base")
def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict]:
    """
    Семантический поиск в базе знаний.
    Возвращает список найденных документов с их содержимым и метаданными.
    """
    return kb.similarity_search(query, k=max_results)


@tool("add_to_knowledge_base")
def add_to_knowledge_base(content: str, title: str) -> str:
    """
    Добавляет новый документ в базу знаний после разбиения на чанки.
    Возвращает сообщение об успешном добавлении.
    """
    docs = split_text_to_documents(content, title)
    kb.add_documents(docs)
    return f"Документ '{title}' успешно добавлен в базу."


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

llm = OllamaLLM(model=LLM_MODEL)

system_prompt = """
Ты – интеллектуальный агент с доступом к локальной базе знаний.
При ответе на запросы используй только информацию из базы, если она есть.
Если необходима дополнительная информация, сначала выполни поиск в базе,
затем сформируй ответ. Не выдавай вымышленных фактов.
"""

tools = [
    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_prompt=system_prompt
)

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


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

def load_documents_from_dir(directory: Path) -> None:
    """
    Загружает все текстовые файлы из указанной папки в базу знаний.
    """
    for file_path in directory.rglob("*"):
        if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}:
            text = file_path.read_text(encoding="utf-8")
            title = file_path.stem
            docs = split_text_to_documents(text, title)
            kb.add_documents(docs)
            print(f"Загружен документ: {title}")


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

def interactive_cli() -> None:
    """
    Простая CLI с командами /add, /search и /quit.
    """
    print("=== RAG‑Агент ===")
    print("Команды:")
    print("/add <title> – добавить новый документ (будет запрошен ввод текста)")
    print("/search <query> – поиск в базе знаний")
    print("/quit – выйти")

    while True:
        try:
            user_input = input("\n> ").strip()
        except EOFError:
            break

        if not user_input:
            continue

        if user_input.lower() == "/quit":
            print("До свидания!")
            break

        if user_input.startswith("/add"):
            parts = user_input.split(maxsplit=1)
            title = parts[1] if len(parts) > 1 else "Untitled"
            print(f"Введите текст для документа '{title}'. Завершите ввод пустой строкой.")
            lines = []
            while True:
                line = input()
                if line == "":
                    break
                lines.append(line)
            content = "\n".join(lines)
            result = add_to_knowledge_base(content, title)
            print(result)

        elif user_input.startswith("/search"):
            query = user_input[len("/search"):].strip()
            if not query:
                print("Введите запрос после /search")
                continue
            results = search_knowledge_base(query, max_results=3)
            if not results:
                print("Ничего не найдено.")
                continue
            for i, res in enumerate(results, 1):
                print(f"\nРезультат {i}:")
                print(f"Титул: {res['metadata'].get('title', 'Неизвестно')}")
                print(res["content"][:500] + ("..." if len(res["content"]) > 500 else ""))

        else:
            # Любой другой ввод считается запросом к агенту
            response = executor.invoke({"input": user_input})
            print(response.get("output", "Нет ответа."))


# --------------------------------------------------------------------------- #
# Точка входа
# --------------------------------------------------------------------------- #

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--init":
        # Инициализация: загрузка документов из папки docs/
        docs_dir = Path("docs")
        if not docs_dir.exists():
            print(f"Папка {docs_dir} не найдена.")
            sys.exit(1)
        load_documents_from_dir(docs_dir)
    else:
        interactive_cli()