#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
solution.py

Основная реализация агента с RAG‑памятью.
Используется стек:
    - FastMCP (FastAPI + OpenAI API wrapper)
    - LangChain для построения цепочки RetrievalQA
    - Qdrant в качестве векторного хранилища
    - rich для красивого вывода

Файл запускается как скрипт и демонстрирует работу агента:
1. Загружает документы из указанной папки.
2. Создаёт эмбеддинги с помощью OpenAI модели.
3. Сохраняет векторные представления в Qdrant.
4. Инициализирует цепочку RetrievalQA.
5. Позволяет задавать вопросы через CLI и выводит ответы.

Требования:
    pip install fastapi uvicorn langchain qdrant-client openai rich
"""

import os
import sys
import argparse
from pathlib import Path
from typing import List

import openai
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores.qdrant import Qdrant
from langchain.chains.question_answering import load_qa_chain
from langchain.llms.openai import OpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter

from rich.console import Console
from rich.prompt import Prompt
from rich.panel import Panel

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

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError(
        "Переменная окружения OPENAI_API_KEY не установлена. "
        "Пожалуйста, получите ключ OpenAI и экспортируйте его."
    )
openai.api_key = OPENAI_API_KEY

# Путь к директории с документами
DOCS_DIR = Path("docs")  # можно изменить при необходимости

# Параметры Qdrant
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
COLLECTION_NAME = "rag_memory"

# --------------------------------------------------------------------------- #
# Утилиты для работы с документами
# --------------------------------------------------------------------------- #

def load_text_files(directory: Path) -> List[str]:
    """
    Загружает все текстовые файлы из указанной директории.
    Возвращает список строк, каждая из которых содержит содержимое одного файла.
    """
    texts = []
    for file_path in directory.glob("**/*.txt"):
        try:
            content = file_path.read_text(encoding="utf-8")
            if content.strip():
                texts.append(content)
        except Exception as exc:
            console.print(f"[red]Не удалось прочитать {file_path}: {exc}[/red]")
    return texts

def split_documents(texts: List[str], chunk_size: int = 1000, overlap: int = 200) -> List[str]:
    """
    Делит большие тексты на чанки фиксированного размера.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=overlap,
        separators=["\n\n", "\n", " ", ""],
    )
    chunks = []
    for text in texts:
        chunks.extend(splitter.split_text(text))
    return chunks

# --------------------------------------------------------------------------- #
# Векторизация и хранение
# --------------------------------------------------------------------------- #

def create_or_load_vectorstore(chunks: List[str]) -> Qdrant:
    """
    Создаёт векторное хранилище Qdrant из списка чанков.
    Если коллекция уже существует, просто загружает её.
    """
    embeddings = OpenAIEmbeddings()
    # Проверяем наличие коллекции
    try:
        vectorstore = Qdrant(
            client=QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT),
            collection_name=COLLECTION_NAME,
            embedding_function=embeddings.embed_query,
        )
        console.print(f"[green]Коллекция '{COLLECTION_NAME}' уже существует. Загружаем...[/green]")
    except Exception:
        # Создаём новую коллекцию
        vectorstore = Qdrant.from_texts(
            texts=chunks,
            embedding=embeddings,
            client=QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT),
            collection_name=COLLECTION_NAME,
        )
        console.print(f"[green]Коллекция '{COLLECTION_NAME}' создана и заполнена.{/green}")
    return vectorstore

# --------------------------------------------------------------------------- #
# Создание цепочки RetrievalQA
# --------------------------------------------------------------------------- #

def build_qa_chain(vectorstore: Qdrant) -> any:
    """
    Инициализирует цепочку RetrievalQA с использованием LLM OpenAI.
    """
    llm = OpenAI(temperature=0.2, model_name="gpt-3.5-turbo")
    chain = load_qa_chain(
        llm,
        chain_type="stuff",
        verbose=False
    )
    return lambda question: chain.run(input_documents=vectorstore.similarity_search(question), question=question)

# --------------------------------------------------------------------------- #
# Основной цикл взаимодействия с пользователем
# --------------------------------------------------------------------------- #

def main():
    console.print(Panel("[bold cyan]RAG‑агент запущен[/bold cyan]", expand=False))

    # 1. Загружаем документы
    if not DOCS_DIR.exists() or not any(DOCS_DIR.glob("**/*.txt")):
        console.print(f"[red]Папка с документами '{DOCS_DIR}' пуста или не существует.[/red]")
        sys.exit(1)

    raw_texts = load_text_files(DOCs_DIR)
    console.print(f"[green]Загружено {len(raw_texts)} файлов. Делим на чанки...[/green]")

    # 2. Разбиваем на чанки
    chunks = split_documents(raw_texts, chunk_size=1000, overlap=200)
    console.print(f"[green]Получено {len(chunks)} чанков.[/green]")

    # 3. Создаём или загружаем векторное хранилище
    vectorstore = create_or_load_vectorstore(chunks)

    # 4. Строим цепочку QA
    qa_function = build_qa_chain(vectorstore)
    console.print("[green]Цепочка RetrievalQA готова![/green]")

    # 5. Взаимодействие с пользователем
    while True:
        try:
            question = Prompt.ask("\n[bold magenta]Введите ваш вопрос (или 'выход' для завершения)[/bold magenta]")
            if question.lower() in {"выход", "exit", "quit"}:
                console.print("[yellow]Завершаем работу агента.[/yellow]")
                break
            answer = qa_function(question)
            console.print(Panel(f"[white]{answer}[/white]", title="Ответ", expand=False))
        except KeyboardInterrupt:
            console.print("\n[yellow]Прерывание пользователем. Завершение...[/yellow]")
            break

if __name__ == "__main__":
    main()