import os
import sys
import logging
from typing import List

from langchain_ollama import ChatOllama
from langchain_qdrant import QdrantStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents 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
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
EMBEDDING_MODEL = os.getenv("NOMIC_EMBED_TEXT_MODEL", "nomic-embed-text")
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")

# Инициализация Ollama и Qdrant
try:
    llm = ChatOllama(model=OLLAMA_MODEL)
except Exception as e:
    logger.error(f"Не удалось подключиться к Ollama: {e}")
    sys.exit(1)

try:
    qdrant_store = QdrantStore(
        url=QDRANT_URL,
        collection_name=COLLECTION_NAME,
        embedding_function=lambda text: llm.embed_query(text),
    )
except Exception as e:
    logger.error(f"Не удалось подключиться к Qdrant: {e}")
    sys.exit(1)

# Текстовый разделитель с гибкой схемой chunking
def split_text(text: str, chunk_size: int = 500, overlap: int = 50) -> List[Document]:
    """
    Делит входной текст на чанки.
    Параметры:
        chunk_size: максимальная длина чанка в символах
        overlap: количество пересекающихся символов между соседними чанками
    Возвращает список Document объектов.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=overlap,
        separators=["\n\n", "\n", " ", ""],
    )
    return [Document(page_content=c) for c in splitter.split_text(text)]

def add_document(file_path: str):
    """Добавляет документ из файла в базу Qdrant."""
    if not os.path.isfile(file_path):
        logger.warning(f"Файл не найден: {file_path}")
        return

    try:
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
    except Exception as e:
        logger.error(f"Ошибка чтения файла {file_path}: {e}")
        return

    docs = split_text(content)
    if not docs:
        logger.warning("Документ пустой после разбиения на чанки.")
        return

    try:
        qdrant_store.add_documents(docs)
        logger.info(f"Добавлено {len(docs)} чанков из файла: {file_path}")
    except Exception as e:
        logger.error(f"Ошибка при добавлении документов в Qdrant: {e}")

def search(query: str, k: int = 5):
    """Ищет релевантные документы и выводит ответы."""
    try:
        docs_and_scores = qdrant_store.similarity_search_with_score(query, k=k)
    except Exception as e:
        logger.error(f"Ошибка при поиске в Qdrant: {e}")
        return

    if not docs_and_scores:
        logger.info("Ничего не найдено.")
        return

    # Формируем контекст из найденных чанков
    context = "\n\n".join([doc.page_content for doc, _ in docs_and_scores])

    prompt = f"""
Ты эксперт по теме. На основе следующего контекста ответь на вопрос:

Контекст:
{context}

Вопрос: {query}
"""

    try:
        response = llm.invoke(prompt)
        print("\nОтвет:\n" + response.content.strip())
    except Exception as e:
        logger.error(f"Ошибка при генерации ответа Ollama: {e}")

def main():
    logger.info("Запущен CLI клиент RAG. Команды: /add <путь>, /search <запрос>, /quit")
    while True:
        try:
            user_input = input("\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nВыход.")
            break

        if not user_input:
            continue

        if user_input.lower() == "/quit":
            logger.info("Завершение работы.")
            break
        elif user_input.startswith("/add"):
            parts = user_input.split(maxsplit=1)
            if len(parts) != 2:
                print("Использование: /add <путь_к_файлу>")
                continue
            add_document(parts[1])
        elif user_input.startswith("/search"):
            parts = user_input.split(maxsplit=1)
            if len(parts) != 2:
                print("Использование: /search <запрос>")
                continue
            search(parts[1])
        else:
            print("Неизвестная команда. Доступные команды: /add, /search, /quit")

if __name__ == "__main__":
    main()