import sys
import logging
from typing import List

# LangChain imports
from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain.schema import Document

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

# Конфигурация Ollama и Qdrant
EMBEDDING_MODEL = "nomic-embed-text"
LLM_MODEL = "llama3"

# Создаём объект для получения эмбеддингов
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)

# Подключаемся к локальному экземпляру Qdrant (по умолчанию localhost:6333)
vector_store = QdrantVectorStore(
    embedding_function=embeddings,
    url="http://localhost:6333",
    collection_name="documents",
)


def middleware(func):
    """
    Декоратор для обработки ошибок и логирования.
    """

    def wrapper(*args, **kwargs):
        try:
            logger.info(f"Выполняется команда: {func.__name__}")
            result = func(*args, **kwargs)
            logger.info("Команда выполнена успешно")
            return result
        except Exception as e:
            logger.exception(f"Ошибка в команде {func.__name__}: {e}")
            print(f"[ERROR] {e}")

    return wrapper


@middleware
def add_document(text: str) -> None:
    """
    Добавляет документ в Qdrant.
    """
    doc = Document(page_content=text)
    vector_store.add_documents([doc])
    print("Документ добавлен.")


@middleware
def search(query: str, k: int = 3) -> List[Document]:
    """
    Выполняет поиск по запросу и выводит результаты.
    """
    results = vector_store.similarity_search_with_score(query, k=k)
    if not results:
        print("Ничего не найдено.")
        return []

    for idx, (doc, score) in enumerate(results, start=1):
        print(f"\nРезультат {idx} (score: {score:.4f})")
        print(doc.page_content[:500] + ("..." if len(doc.page_content) > 500 else ""))
    return [doc for doc, _ in results]


def parse_command(line: str):
    """
    Разбирает строку команды и вызывает соответствующую функцию.
    """
    if not line.strip():
        return

    parts = line.split(maxsplit=1)
    cmd = parts[0].lower()

    if cmd == "/add":
        if len(parts) < 2:
            print("[ERROR] Необходимо указать текст для добавления.")
            return
        add_document(parts[1])
    elif cmd == "/search":
        if len(parts) < 2:
            print("[ERROR] Необходимо указать запрос для поиска.")
            return
        search(parts[1])
    elif cmd == "/quit":
        print("Выход из программы.")
        sys.exit(0)
    else:
        print(f"[ERROR] Нераспознанная команда: {cmd}")


def main():
    """
    Основной цикл CLI.
    """
    logger.info("Запуск клиента Human-in-the-Loop")
    print("Добро пожаловать! Используйте команды /add, /search и /quit.")
    while True:
        try:
            line = input("> ")
            parse_command(line)
        except KeyboardInterrupt:
            print("\nПрервано пользователем. Выход.")
            break
        except EOFError:
            print("\nEOF. Выход.")
            break


# -------------------- UNIT TESTS --------------------
import unittest
from unittest.mock import patch, MagicMock

class TestClient(unittest.TestCase):
    def setUp(self):
        # Подменяем vector_store для тестов
        self.original_vector_store = vector_store
        mock_vs = MagicMock()
        mock_vs.add_documents.return_value = None
        mock_vs.similarity_search_with_score.return_value = [
            (Document(page_content="Test content 1"), 0.9),
            (Document(page_content="Test content 2"), 0.8),
        ]
        global vector_store
        vector_store = mock_vs

    def tearDown(self):
        global vector_store
        vector_store = self.original_vector_store

    @patch('builtins.print')
    def test_add_document_success(self, mock_print):
        add_document("Sample text")
        mock_print.assert_called_with("Документ добавлен.")

    @patch('builtins.print')
    def test_search_results(self, mock_print):
        results = search("query")
        self.assertEqual(len(results), 2)
        # Проверяем, что print был вызван хотя бы один раз
        self.assertTrue(mock_print.call_count > 0)

    @patch('builtins.print')
    def test_parse_unknown_command(self, mock_print):
        parse_command("/unknown")
        mock_print.assert_called_with("[ERROR] Нераспознанная команда: /unknown")

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "test":
        # Запуск тестов
        unittest.main(argv=[sys.argv[0]])
    else:
        main()