#!/usr/bin/env python3
"""
Script de nettoyage pour PrivateGPT

Ce script permet de lister et supprimer les documents ingérés dans PrivateGPT.
Options:
- Lister tous les documents
- Supprimer des documents par préfixe (ex: "temp_section_")
- Supprimer des documents par motif
- Supprimer tous les documents

Usage:
    python nettoyer_pgpt.py --list
    python nettoyer_pgpt.py --delete-prefix "temp_section_"
    python nettoyer_pgpt.py --delete-pattern "rapport_.*\.md"
    python nettoyer_pgpt.py --delete-all
"""

import argparse
import json
import re
import requests
import sys
import time
import uuid
from typing import List, Dict, Any, Optional

# Configuration de l'API PrivateGPT
PGPT_URL = "http://127.0.0.1:8001"
API_URL = f"{PGPT_URL}/v1"


def check_api_availability() -> bool:
    """Vérifie si l'API PrivateGPT est disponible"""
    try:
        response = requests.get(f"{PGPT_URL}/health")
        if response.status_code == 200:
            print("✅ API PrivateGPT disponible")
            return True
        else:
            print(f"❌ L'API PrivateGPT a retourné le code d'état {response.status_code}")
            return False
    except requests.RequestException as e:
        print(f"❌ Erreur de connexion à l'API PrivateGPT: {e}")
        return False


def list_documents() -> List[Dict[str, Any]]:
    """Liste tous les documents ingérés et renvoie la liste des métadonnées"""
    try:
        # Récupérer la liste des documents
        response = requests.get(f"{API_URL}/ingest/list")
        response.raise_for_status()
        data = response.json()
        
        # Format de réponse OpenAI
        if "data" in data:
            documents = data.get("data", [])
        # Format alternatif
        else:
            documents = data.get("documents", [])
        
        # Construire une liste normalisée des documents
        normalized_docs = []
        for doc in documents:
            doc_id = doc.get("doc_id") or doc.get("id")
            metadata = doc.get("doc_metadata", {})
            filename = metadata.get("file_name") or metadata.get("filename", "Inconnu")
            
            normalized_docs.append({
                "id": doc_id,
                "filename": filename,
                "metadata": metadata
            })
            
        return normalized_docs
    
    except Exception as e:
        print(f"❌ Erreur lors de la récupération des documents: {e}")
        return []


def print_documents(documents: List[Dict[str, Any]]) -> None:
    """Affiche la liste des documents de façon lisible"""
    if not documents:
        print("📋 Aucun document trouvé dans PrivateGPT")
        return
    
    print(f"📋 {len(documents)} documents trouvés dans PrivateGPT:")
    
    # Regrouper par nom de fichier pour un affichage plus compact
    files_grouped = {}
    for doc in documents:
        filename = doc["filename"]
        if filename not in files_grouped:
            files_grouped[filename] = []
        files_grouped[filename].append(doc["id"])
    
    # Afficher les résultats groupés
    for i, (filename, ids) in enumerate(files_grouped.items(), 1):
        print(f"{i}. {filename} ({len(ids)} chunks)")
        if args.verbose:
            for j, doc_id in enumerate(ids, 1):
                print(f"   {j}. ID: {doc_id}")


def delete_document(doc_id: str) -> bool:
    """Supprime un document par son ID"""
    try:
        response = requests.delete(f"{API_URL}/ingest/{doc_id}")
        if response.status_code == 200:
            return True
        else:
            print(f"⚠️ Échec de la suppression de l'ID {doc_id}: Code {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Erreur lors de la suppression de l'ID {doc_id}: {e}")
        return False


def delete_documents_by_criteria(documents: List[Dict[str, Any]], 
                                prefix: Optional[str] = None, 
                                pattern: Optional[str] = None,
                                delete_all: bool = False) -> int:
    """
    Supprime des documents selon différents critères
    Retourne le nombre de documents supprimés
    """
    if not documents:
        print("❌ Aucun document à supprimer")
        return 0
    
    if not (prefix or pattern or delete_all):
        print("❌ Aucun critère de suppression spécifié")
        return 0
    
    # Comptage des suppressions réussies
    success_count = 0
    
    # Filtrer les documents à supprimer
    docs_to_delete = []
    
    if delete_all:
        docs_to_delete = documents
        print(f"🗑️ Suppression de tous les documents ({len(documents)} chunks)...")
    elif prefix:
        docs_to_delete = [doc for doc in documents if doc["filename"].startswith(prefix)]
        print(f"🗑️ Suppression des documents dont le nom commence par '{prefix}' ({len(docs_to_delete)} chunks)...")
    elif pattern:
        try:
            regex = re.compile(pattern)
            docs_to_delete = [doc for doc in documents if regex.search(doc["filename"])]
            print(f"🗑️ Suppression des documents correspondant au motif '{pattern}' ({len(docs_to_delete)} chunks)...")
        except re.error as e:
            print(f"❌ Expression régulière invalide: {e}")
            return 0
    
    # Demander confirmation si beaucoup de documents
    if len(docs_to_delete) > 5 and not args.force:
        confirm = input(f"⚠️ Vous êtes sur le point de supprimer {len(docs_to_delete)} chunks. Confirmer ? (o/N) ")
        if confirm.lower() != 'o':
            print("❌ Opération annulée")
            return 0
    
    # Supprimer les documents
    for doc in docs_to_delete:
        if delete_document(doc["id"]):
            success_count += 1
            if args.verbose:
                print(f"✅ Document supprimé: {doc['filename']} (ID: {doc['id']})")
        
        # Petite pause pour éviter de surcharger l'API
        time.sleep(0.1)
    
    print(f"✅ {success_count}/{len(docs_to_delete)} documents supprimés avec succès")
    return success_count


def generate_unique_prefix() -> str:
    """Génère un préfixe unique basé sur un UUID pour différencier les fichiers temporaires"""
    unique_id = str(uuid.uuid4())[:8]  # Prendre les 8 premiers caractères de l'UUID
    return f"temp_{unique_id}_"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Utilitaire de nettoyage pour PrivateGPT")
    
    # Options principales
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--list", action="store_true", help="Lister tous les documents ingérés")
    group.add_argument("--delete-prefix", type=str, help="Supprimer les documents dont le nom commence par PREFIX")
    group.add_argument("--delete-pattern", type=str, help="Supprimer les documents dont le nom correspond au motif PATTERN (regex)")
    group.add_argument("--delete-all", action="store_true", help="Supprimer tous les documents (⚠️ DANGER)")
    group.add_argument("--generate-prefix", action="store_true", help="Générer un préfixe unique pour les fichiers temporaires")
    
    # Options additionnelles
    parser.add_argument("--force", action="store_true", help="Ne pas demander de confirmation")
    parser.add_argument("--verbose", action="store_true", help="Afficher plus de détails")
    
    args = parser.parse_args()
    
    # Vérifier la disponibilité de l'API
    if not check_api_availability():
        sys.exit(1)
    
    # Générer un préfixe unique
    if args.generate_prefix:
        unique_prefix = generate_unique_prefix()
        print(f"🔑 Préfixe unique généré: {unique_prefix}")
        print(f"Utilisez ce préfixe pour les fichiers temporaires de votre script.")
        sys.exit(0)
    
    # Récupérer la liste des documents
    documents = list_documents()
    
    # Traiter selon l'option choisie
    if args.list:
        print_documents(documents)
    elif args.delete_prefix:
        delete_documents_by_criteria(documents, prefix=args.delete_prefix)
    elif args.delete_pattern:
        delete_documents_by_criteria(documents, pattern=args.delete_pattern)
    elif args.delete_all:
        delete_documents_by_criteria(documents, delete_all=True)