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

"""
solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama.
"""

import argparse
from pathlib import Path
from typing import List, Dict

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

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

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

EMBEDDING_MODEL = "nomic-embed-text"   # Ollama embedding model
LLM_MODEL = "ollama:llama3"            # Ollama LLM (формат для create_agent)

# --------------------------------------------------------------------------- #
# Векторное хранилище и вспомогательные функции
# --------------------------------------------------------------------------- #

class KnowledgeBase:
    """
    Обёртка над QdrantVectorStore с поддержкой чанкинга и инструментов.
    """

    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(
            embedding=self.embeddings,
            url=f"http://{host}:{port}",
            collection_name=collection_name,
        )
        # Чанкер для разбивки документов
        self.splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)

    def add_document(self, content: str, title: str) -> None:
        """
        Добавляет документ в базу после разбиения на чанки.
        """
        chunks = self.splitter.split_text(content)
        docs = [
            Document(
                page_content=chunk,
                metadata={"title": title, "source": f"{title} (chunk {i})"}
            )
            for i, chunk in enumerate(chunks)
        ]
        self.store.add_documents(docs)

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


# --------------------------------------------------------------------------- #
# Глобальный экземпляр базы знаний
# --------------------------------------------------------------------------- #

kb = KnowledgeBase()

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

@tool
def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """
    Поиск в базе знаний.
    Возвращает строку с найденными результатами в формате Markdown.
    """
    results = kb.search(query, max_results)
    if not results:
        return "Ничего не найдено."
    lines = []
    for r in results:
        lines.append(
            f"**{r['title']}** ({r['source']}) – {r['score']:.2f}\n> {r['content'][:200]}..."
        )
    return "\n\n".join(lines)


@tool
def add_to_knowledge_base(content: str, title: str) -> str:
    """
    Добавление нового документа в базу знаний.
    """
    kb.add_document(content, title)
    return f"Документ «{title}» успешно добавлен."


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

def build_agent():
    """
    Возвращает готовый агент с инструментами RAG‑памяти.
    Используется актуальный API: create_agent(model, tools, system_prompt).
    """
    system_prompt = (
        "Ты – интеллектуальный ассистент, который использует локальную базу знаний. "
        "При ответе на запросы сначала ищи релевантную информацию в базе через инструмент "
        "`search_knowledge_base`. Если нужно сохранить новую информацию, добавляй её через "
        "`add_to_knowledge_base`."
    )

    agent = create_agent(
        LLM_MODEL,
        tools=[search_knowledge_base, add_to_knowledge_base],
        system_prompt=system_prompt,
    )
    return agent


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

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


# --------------------------------------------------------------------------- #
# CLI клиент
# --------------------------------------------------------------------------- #

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

    while True:
        try:
            inp = input("\n> ").strip()
        except EOFError:
            break
        if not inp:
            continue

        if inp.lower() == "/quit":
            break

        elif inp.startswith("/add "):
            title = inp[5:].strip()
            print("Введите содержимое документа (завершите ввод Ctrl+D / Ctrl+Z):")
            try:
                content_lines = []
                while True:
                    line = input()
                    content_lines.append(line)
            except EOFError:
                pass
            content = "\n".join(content_lines)
            result = add_to_knowledge_base.invoke({"content": content, "title": title})
            print(result)

        elif inp.startswith("/search "):
            query = inp[8:].strip()
            results = kb.search(query, max_results=5)
            if not results:
                print("Ничего не найдено.")
            else:
                for r in results:
                    print(f"\n**{r['title']}** ({r['source']}) – {r['score']:.2f}")
                    print(r["content"][:400] + "...")

        else:
            # Любой другой ввод – запрос к агенту
            # Актуальный API: invoke принимает {"messages": [...]}
            result = agent.invoke(
                {"messages": [{"role": "user", "content": inp}]}
            )
            # Извлекаем текст последнего сообщения агента
            messages = result.get("messages", [])
            if messages:
                print(messages[-1].content)


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

def main() -> None:
    parser = argparse.ArgumentParser(description="Агент с RAG‑памятью")
    parser.add_argument(
        "--load-dir",
        type=str,
        help="Путь к директории с документами для предварительной загрузки"
    )
    args = parser.parse_args()

    if args.load_dir:
        load_documents_from_dir(Path(args.load_dir))

    agent = build_agent()
    interactive_cli(agent)


if __name__ == "__main__":
    main()