# rag_agent.py
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_react_agent, AgentExecutor
from langchain.prompts import PromptTemplate
from langchain.schema import Document
from qdrant_client import QdrantClient as _QdrantClient
from qdrant_client.http.models import Distance, VectorParams

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

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 = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
LLM_MODEL = os.getenv("LLM_MODEL", "llama3")

# ---------------------------------------------------------------------------
# Векторное хранилище (Qdrant + Ollama-эмбеддинги)
# ---------------------------------------------------------------------------

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

    def __init__(self, collection_name: str = COLLECTION_NAME):
        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)

        # Низкоуровневый клиент для управления коллекцией
        self._client = _QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)

        # Создаём коллекцию, если её ещё нет
        if not self._collection_exists(collection_name):
            self._client.create_collection(
                collection_name=collection_name,
                vectors_config=VectorParams(size=768, distance=Distance.COSINE),
            )

        self.store = QdrantVectorStore(
            client=self._client,
            collection_name=collection_name,
            embedding=self.embeddings,
        )

    def _collection_exists(self, name: str) -> bool:
        try:
            self._client.get_collection(name)
            return True
        except Exception:
            return False

    def add_documents(self, docs: List[Document]) -> None:
        self.store.add_documents(docs)

    def search(self, query: str, k: int = 5) -> List[tuple]:
        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]:
    """Разбивает текст на чанки с метаданными."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        separators=["\n\n", "\n", " ", ""],
    )
    return [
        Document(page_content=chunk, metadata={"title": title})
        for chunk in splitter.split_text(text)
    ]


# ---------------------------------------------------------------------------
# Инициализация хранилища (глобальный синглтон для инструментов)
# ---------------------------------------------------------------------------

rag_store = RagVectorStore()


# ---------------------------------------------------------------------------
# LangChain-инструменты агента (@tool)
# ---------------------------------------------------------------------------

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

    Args:
        query: поисковый запрос.
        max_results: максимальное количество результатов.

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


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

    Args:
        content: полный текст документа.
        title: заголовок документа.

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


# ---------------------------------------------------------------------------
# Агент на базе LangChain (create_react_agent)
# ---------------------------------------------------------------------------

llm = OllamaLLM(model=LLM_MODEL)

tools = [search_knowledge_base, add_to_knowledge_base]

prompt = PromptTemplate.from_template(
    "Вы — интеллектуальный агент с доступом к локальной базе знаний.\n"
    "При необходимости используйте инструменты search_knowledge_base и add_to_knowledge_base.\n\n"
    "Доступные инструменты:\n{tools}\n\n"
    "Формат ответа (строго соблюдайте):\n"
    "Question: вопрос пользователя\n"
    "Thought: что нужно сделать\n"
    "Action: инструмент из [{tool_names}]\n"
    "Action Input: входные данные для инструмента\n"
    "Observation: результат инструмента\n"
    "... (можно повторять Thought/Action/Observation)\n"
    "Thought: теперь я знаю ответ\n"
    "Final Answer: итоговый ответ пользователю\n\n"
    "Begin!\n\n"
    "Question: {input}\n"
    "Thought: {agent_scratchpad}"
)

agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)


# ---------------------------------------------------------------------------
# Загрузка документов из директории (для init_loader.py)
# ---------------------------------------------------------------------------

def load_documents_from_directory(directory: str) -> None:
    """Загружает все .txt и .md файлы из папки в базу знаний."""
    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")
            add_to_knowledge_base.invoke({"content": text, "title": file_path.stem})
            print(f"Загружен: {file_path.name}")


# ---------------------------------------------------------------------------
# Запуск напрямую — интерактивный CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    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 "):
            file_path = Path(user_input[5:].strip())
            if file_path.is_file():
                text = file_path.read_text(encoding="utf-8")
                print(add_to_knowledge_base.invoke({"content": text, "title": file_path.stem}))
            else:
                print(f"Файл не найден: {file_path}")
        elif user_input.startswith("/search "):
            query = user_input[8:].strip()
            for r in search_knowledge_base.invoke({"query": query, "max_results": 3}):
                print(r)
        else:
            response = agent_executor.invoke({"input": user_input})
            print(response["output"])