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

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

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

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.agents import create_agent, AgentExecutor, Tool
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 = "llama3"                   # Ollama LLM

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

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)
        ]
        # Добавляем в Qdrant
        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:
            out = {
                "score": score,
                "title": doc.metadata.get("title", ""),
                "source": doc.metadata.get("source", ""),
                "content": doc.page_content
            }
            output.append(out)
        return output


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

kb = KnowledgeBase()

@tool
def search_knowledge_base(query: str, max_results: int) -> 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() -> AgentExecutor:
    """
    Возвращает готовый агент с инструментами RAG‑памяти.
    """
    system_prompt = (
        "Ты – интеллектуальный ассистент, который использует локальную базу знаний. "
        "При ответе на запросы сначала ищи релевантную информацию в базе через tool "
        "`search_knowledge_base`. Если нужна новая информация, добавляй её через "
        "`add_to_knowledge_base`."
    )
    tools = [
        Tool.from_function(
            func=search_knowledge_base,
            name="search_knowledge_base",
            description="Семантический поиск в базе знаний. Вход: query, max_results. Выход: Markdown с результатами."
        ),
        Tool.from_function(
            func=add_to_knowledge_base,
            name="add_to_knowledge_base",
            description="Добавить новый документ в базу знаний. Вход: content, title. Выход: подтверждение."
        )
    ]
    llm = ChatOllama(model=LLM_MODEL)
    agent = create_agent(
        llm=llm,
        tools=tools,
        system_prompt=system_prompt,
        verbose=True
    )
    return AgentExecutor(agent=agent, tools=tools, verbose=True)


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

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_executor: AgentExecutor) -> 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

        if inp.startswith("/add "):
            title = inp[5:].strip()
            print("Введите содержимое документа (конец строки с EOF Ctrl+D):")
            try:
                content_lines = []
                while True:
                    line = input()
                    content_lines.append(line)
            except EOFError:
                pass
            content = "\n".join(content_lines)
            result = add_to_knowledge_base(content, 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:
            # Любой другой ввод – запрос к агенту
            response = agent_executor.invoke({"input": inp})
            print(response.get("output", ""))


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

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_executor = build_agent()
    interactive_cli(agent_executor)


if __name__ == "__main__":
    main()