#!/usr/bin/env python3
"""
Démonstration de RAG (Retrieval-Augmented Generation) avec LlamaIndex.
Supporte les fichiers PDF, TXT et MD.
"""

import os
import shutil
from pathlib import Path
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    load_index_from_storage,
    Settings,
)
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding

# --- Configuration ---
PERSIST_DIR = "./storage"
DOCS_DIR = "./documents"
SUPPORTED_EXTENSIONS = {".pdf", ".txt", ".md"}
OLLAMA_LLM = "qwen2.5:1.5b"
OLLAMA_EMBED = "nomic-embed-text"
SHOW_CHUNKS = False

Settings.llm = Ollama(model=OLLAMA_LLM, request_timeout=120)
Settings.embed_model = OllamaEmbedding(model_name=OLLAMA_EMBED)


def ensure_dirs():
    """Crée les répertoires nécessaires s'ils n'existent pas."""
    Path(DOCS_DIR).mkdir(exist_ok=True)
    Path(PERSIST_DIR).mkdir(exist_ok=True)


def get_index():
    """Charge l'index existant ou en crée un nouveau."""
    if os.path.exists(PERSIST_DIR) and os.listdir(PERSIST_DIR):
        storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
        return load_index_from_storage(storage_context)
    return VectorStoreIndex([], embed_model=Settings.embed_model)


def save_index(index):
    """Sauvegarde l'index sur le disque."""
    index.storage_context.persist(persist_dir=PERSIST_DIR)


def list_documents():
    """Liste les documents disponibles dans le répertoire docs."""
    files = [f for f in Path(DOCS_DIR).iterdir() if f.suffix.lower() in SUPPORTED_EXTENSIONS]
    if not files:
        print("  Aucun document trouvé.")
        return []
    for i, f in enumerate(files, 1):
        size = f.stat().st_size
        print(f"  {i}. {f.name} ({size} octets)")
    return files


def add_document(index):
    """Ajoute un document au répertoire et met à jour l'index."""
    path = input("Chemin du document à ajouter : ").strip()
    if not path:
        print("  Chemin invalide.")
        return index

    src = Path(path)
    if not src.exists():
        print(f"  Fichier introuvable : {path}")
        return index

    if src.suffix.lower() not in SUPPORTED_EXTENSIONS:
        print(f"  Extension non supportée : {src.suffix}")
        print(f"  Formats acceptés : {', '.join(SUPPORTED_EXTENSIONS)}")
        return index

    dest = Path(DOCS_DIR) / src.name
    if dest.exists():
        print(f"  Un fichier '{src.name}' existe déjà dans les documents.")
        overwrite = input("  Écraser ? (o/n) : ").strip().lower()
        if overwrite != "o":
            return index

    shutil.copy2(src, dest)
    print(f"  Document copié : {src.name}")

    documents = SimpleDirectoryReader(DOCS_DIR).load_data()
    index = VectorStoreIndex.from_documents(documents)
    save_index(index)
    print("  Index mis à jour.")
    return index


def delete_document(index):
    """Supprime un document et reconstruit l'index."""
    files = list_documents()
    if not files:
        return index

    try:
        choice = int(input("Numéro du document à supprimer : ")) - 1
        if choice < 0 or choice >= len(files):
            print("  Numéro invalide.")
            return index
    except ValueError:
        print("  Entrée invalide.")
        return index

    target = files[choice]
    confirm = input(f"  Supprimer '{target.name}' ? (o/n) : ").strip().lower()
    if confirm != "o":
        print("  Annulé.")
        return index

    target.unlink()
    print(f"  Document supprimé : {target.name}")

    remaining = [f for f in Path(DOCS_DIR).iterdir() if f.suffix.lower() in SUPPORTED_EXTENSIONS]
    if remaining:
        documents = SimpleDirectoryReader(DOCS_DIR).load_data()
        index = VectorStoreIndex.from_documents(documents)
    else:
        index = VectorStoreIndex([])
    save_index(index)
    print("  Index mis à jour.")
    return index


def search(index):
    """Pose une question à l'index RAG."""
    if not os.listdir(DOCS_DIR):
        print("  Aucun document indexé. Ajoutez d'abord un document.")
        return

    query = input("Votre question : ").strip()
    if not query:
        print("  Question vide.")
        return

    query_engine = index.as_query_engine(llm=Settings.llm)
    response = query_engine.query(query)
    print(f"\n  Réponse : \033[92m{response.response}\033[0m")

    if response.source_nodes:
        print("\n  Sources :")
        for node in response.source_nodes:
            source = node.metadata.get("file_name", "inconnu")
            score = node.score if node.score else "N/A"
            print(f"    - {source} (score: {score})")
            if SHOW_CHUNKS:
                text = node.node.get_text()
                print(f"      Chunk : \033[93m{text}\033[0m")


def toggle_chunks():
    """Active ou désactive l'affichage des chunks."""
    global SHOW_CHUNKS
    SHOW_CHUNKS = not SHOW_CHUNKS
    status = "activé" if SHOW_CHUNKS else "désactivé"
    print(f"  Affichage des chunks {status}.")


def main():
    ensure_dirs()
    index = get_index()

    menu = """
╔══════════════════════════════════╗
║     RAG - LlamaIndex Demo        ║
╠══════════════════════════════════╣
║  1. Ajouter un document          ║
║  2. Supprimer un document        ║
║  3. Rechercher                   ║
║  4. Afficher les documents       ║
║  5. Afficher les chunks : OFF    ║
║  6. Quitter                      ║
╚══════════════════════════════════╝"""

    actions = {
        "1": add_document,
        "2": delete_document,
        "4": list_documents,
    }

    while True:
        status = "ON " if SHOW_CHUNKS else "OFF"
        print(f"""
╔══════════════════════════════════╗
║     RAG - LlamaIndex Demo        ║
╠══════════════════════════════════╣
║  1. Ajouter un document          ║
║  2. Supprimer un document        ║
║  3. Rechercher                   ║
║  4. Afficher les documents       ║
║  5. Afficher les chunks : {status}    ║
║  6. Quitter                      ║
╚══════════════════════════════════╝""")
        choice = input("\n  Votre choix : ").strip()

        if choice == "6":
            print("  Au revoir !")
            break

        if choice == "3":
            search(index)
        elif choice == "4":
            list_documents()
        elif choice == "5":
            toggle_chunks()
        elif choice in actions:
            index = actions[choice](index)
        else:
            print("  Choix invalide.")


if __name__ == "__main__":
    main()