import os
import logging
from typing import Any, Dict, List

from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_qdrant import QdrantStore
from langchain_core.documents import Document
from langchain_core.runnables import RunnableLambda

# Настройка логирования
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# Конфигурация подключения к Qdrant
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "knowledge_base")

# Инициализация эмбеддингов и модели Ollama
try:
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
except Exception as e:
    logger.exception("Не удалось подключиться к Ollama Embeddings: %s", e)
    raise

try:
    llm = ChatOllama(model="llama3")
except Exception as e:
    logger.exception("Не удалось подключиться к Ollama LLM: %s", e)
    raise

# Инициализация хранилища Qdrant
try:
    qdrant_store = QdrantStore(
        url=QDRANT_URL,
        collection_name=COLLECTION_NAME,
        embeddings=embeddings,
    )
except Exception as e:
    logger.exception("Не удалось подключиться к Qdrant: %s", e)
    raise


# ==========================
# Middleware для обработки ошибок и логирования
# ==========================
def middleware(func):
    """
    Декоратор, который оборачивает функцию в try/except,
    логирует входные параметры и результат.
    """

    def wrapper(*args, **kwargs):
        func_name = func.__name__
        logger.debug("Вызов %s с args=%s kwargs=%s", func_name, args, kwargs)
        try:
            result = func(*args, **kwargs)
            logger.debug("Результат %s: %s", func_name, result)
            return result
        except Exception as exc:
            logger.exception("Ошибка в %s: %s", func_name, exc)
            raise

    return wrapper


# ==========================
# RAG инструменты
# ==========================
@middleware
def search_knowledge_base(query: str) -> List[Document]:
    """
    Поиск релевантных документов в Qdrant.
    Возвращает список Document с полями content и metadata.
    """
    if not query:
        raise ValueError("Запрос не может быть пустым")

    try:
        docs = qdrant_store.similarity_search(query, k=5)
        logger.info("Найдено %d документов по запросу '%s'", len(docs), query)
        return docs
    except Exception as exc:
        logger.exception("Ошибка при поиске: %s", exc)
        raise


@middleware
def add_to_knowledge_base(content: str, metadata: Dict[str, Any] | None = None) -> None:
    """
    Добавляет новый документ в Qdrant.
    content – текст документа,
    metadata – произвольные метаданные (может быть None).
    """
    if not content:
        raise ValueError("Содержимое не может быть пустым")

    try:
        doc = Document(page_content=content, metadata=metadata or {})
        qdrant_store.add_documents([doc])
        logger.info("Документ добавлен в коллекцию '%s'", COLLECTION_NAME)
    except Exception as exc:
        logger.exception("Ошибка при добавлении документа: %s", exc)
        raise


# ==========================
# Пример использования (можно использовать как тесты)
# ==========================
if __name__ == "__main__":
    # Минимальный набор unit‑тестов
    import unittest

    class TestAgent(unittest.TestCase):
        def test_add_and_search(self):
            test_content = "Тестовый документ для проверки RAG."
            test_metadata = {"source": "unit_test"}

            # Добавляем документ
            add_to_knowledge_base(test_content, test_metadata)

            # Ищем по ключевому слову
            results = search_knowledge_base("тестовый")
            self.assertTrue(any(test_content in doc.page_content for doc in results))

        def test_empty_query(self):
            with self.assertRaises(ValueError):
                search_knowledge_base("")

        def test_empty_content(self):
            with self.assertRaises(ValueError):
                add_to_knowledge_base("")

    unittest.main(argv=["first-arg-is-ignored"], exit=False)