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_agent, AgentExecutor
from langchain.schema import Document
from rag_tools import search_knowledge_base, add_to_knowledge_base

# --------------------------------------------------------------------------- #
# 1. Векторное хранилище на базе Qdrant
# --------------------------------------------------------------------------- #

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

    def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333):
        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
        self.store = QdrantVectorStore(
            embeddings=self.embeddings,
            url=f"http://{host}:{port}",
            collection_name=collection_name,
        )

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

    def search(self, query: str, limit: int = 5) -> List[Document]:
        """Поиск по запросу с ограничением на количество результатов."""
        return self.store.similarity_search(query=query, k=limit)


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

def chunk_text(text: str, title: str = "") -> List[Document]:
    """
    Делит текст на чанки и возвращает список Document с метаданными.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        length_function=len
    )
    chunks = splitter.split_text(text)
    docs = [
        Document(page_content=chunk, metadata={"title": title})
        for chunk in chunks
    ]
    return docs


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

kb = KnowledgeBase()

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


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

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

LLM_MODEL = "llama3"
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: {input}\n"
    "Thought: {agent_scratchpad}\n"
    "Action: действие из [{tool_names}]\n"
    "Action Input: входные данные\n"
    "Observation: результат\n"
    "Final Answer: итоговый ответ"
)

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

# --------------------------------------------------------------------------- #
# 5. Инициализация из директории (для загрузки файлов)
# --------------------------------------------------------------------------- #

def load_documents_from_dir(directory: str) -> None:
    """
    Загружает все .txt файлы из указанной папки в базу знаний.
    """
    for file_path in Path(directory).glob("*.txt"):
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
        title = file_path.stem
        add_to_knowledge_base(content=content, title=title)

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

def interactive_cli() -> None:
    """
    Простая CLI с командами /add, /search и /quit.
    """
    init_agent()
    print("=== RAG-агент готов к работе ===")
    while True:
        try:
            user_input = input("\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            break

        if not user_input:
            continue

        if user_input.lower() == "/quit":
            print("Завершение работы.")
            break
        elif user_input.startswith("/add"):
            # /add путь_к_файлу
            _, path = user_input.split(maxsplit=1)
            with open(path, "r", encoding="utf-8") as f:
                content = f.read()
            title = Path(path).stem
            print(add_to_knowledge_base(content, title))
        elif user_input.startswith("/search"):
            # /search запрос
            _, query = user_input.split(maxsplit=1)
            results = search_knowledge_base(query, max_results=3)
            for i, res in enumerate(results, 1):
                print(f"\n--- Результат {i} ---\n{res}")
        else:
            # Любой другой ввод считается запросом к агенту
            response = agent_executor.invoke({"input": user_input})
            print(response["output"])

if __name__ == "__main__":
    interactive_cli()