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

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

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

from langchain_ollama import OllamaEmbeddings
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

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

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

EMBEDDING_MODEL  = "nomic-embed-text"
LLM_MODEL        = "ollama:llama3"   # формат model-string для create_agent

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


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

class KnowledgeBase:
    """
    Инкапсулирует работу с Qdrant и эмбеддингами Ollama.

    Инициализация через QdrantVectorStore.from_existing_collection —
    подключается к уже существующей коллекции (или создаёт при первом
    add_documents, когда Qdrant сам управляет коллекцией через URL).
    """

    def __init__(self,
                 host: str = QDRANT_HOST,
                 port: int = QDRANT_PORT,
                 collection_name: str = COLLECTION_NAME):
        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
        # Подключаемся к существующей коллекции (или используем URL-режим,
        # при котором QdrantVectorStore создаёт коллекцию автоматически).
        self.store = QdrantVectorStore.from_existing_collection(
            embedding=self.embeddings,
            collection_name=collection_name,
            url=f"http://{host}:{port}",
        )

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

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


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

def split_text_to_documents(text: str, title: str) -> List[Document]:
    """
    Разбивает текст на чанки через RecursiveCharacterTextSplitter.
    Каждый чанк получает метаданные: title и порядковый номер чанка.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=MAX_CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
    )
    chunks = splitter.split_text(text)
    return [
        Document(
            page_content=chunk,
            metadata={"title": title, "chunk_index": i},
        )
        for i, chunk in enumerate(chunks)
    ]


# --------------------------------------------------------------------------- #
# Глобальный экземпляр базы знаний + инструменты агента (@tool)
# --------------------------------------------------------------------------- #

kb = KnowledgeBase()


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


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


# --------------------------------------------------------------------------- #
# Агент с RAG‑интеграцией
# --------------------------------------------------------------------------- #

system_prompt = (
    "Ты – интеллектуальный агент с доступом к локальной базе знаний.\n"
    "При ответе на запросы сначала ищи релевантную информацию через инструмент "
    "`search_knowledge_base`. Если нужно сохранить новую информацию — используй "
    "`add_to_knowledge_base`. Не выдумывай факты, которых нет в базе."
)

# create_agent(model, tools, system_prompt) — актуальный API LangChain.
# Первый аргумент — строка модели в формате "provider:model".
# Возвращает граф-агент с .invoke() / .stream() — AgentExecutor не нужен.
agent = create_agent(
    LLM_MODEL,
    tools=[search_knowledge_base, add_to_knowledge_base],
    system_prompt=system_prompt,
)


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

def load_documents_from_dir(directory: Path) -> None:
    """
    Загружает все .txt и .md файлы из указанной папки в базу знаний.
    """
    loaded = 0
    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"Загружен: {file_path}  ({len(docs)} чанков)")
            loaded += 1
    print(f"\nИтого загружено файлов: {loaded}")


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

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

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

        if not user_input:
            continue

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

        # /add <title>
        elif user_input.startswith("/add"):
            parts = user_input.split(maxsplit=1)
            title = parts[1].strip() if len(parts) > 1 else "Untitled"
            print(f"Вводите текст для «{title}». Пустая строка — конец ввода.")
            lines: List[str] = []
            while True:
                line = input()
                if line == "":
                    break
                lines.append(line)
            content = "\n".join(lines)
            # Вызов @tool через .invoke() — правильный способ вызова инструмента
            result = add_to_knowledge_base.invoke({"content": content, "title": title})
            print(result)

        # /search <query>
        elif user_input.startswith("/search"):
            query = user_input[len("/search"):].strip()
            if not query:
                print("Укажите запрос: /search <query>")
                continue
            results = search_knowledge_base.invoke({"query": query, "max_results": 3})
            if not results:
                print("Ничего не найдено.")
            else:
                for i, res in enumerate(results, 1):
                    title = res["metadata"].get("title", "—")
                    score = res.get("score", 0)
                    print(f"\nРезультат {i} | {title} | score={score:.4f}")
                    snippet = res["content"]
                    print(snippet[:500] + ("..." if len(snippet) > 500 else ""))

        # Запрос к агенту
        else:
            # Актуальный API: invoke принимает {"messages": [...]},
            # ответ — в result["messages"][-1].content
            result   = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
            messages = result.get("messages", [])
            if messages:
                print(messages[-1].content)
            else:
                print("Нет ответа.")


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

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--init":
        # python solution.py --init   →  загрузить документы из ./docs/
        docs_dir = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("docs")
        if not docs_dir.exists():
            print(f"Папка '{docs_dir}' не найдена.")
            sys.exit(1)
        load_documents_from_dir(docs_dir)
    else:
        interactive_cli()