#!/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
"""

import json
import re
import requests
import time
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 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 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(pattern) -> int:
    """
    Supprime des documents selon différents critères
    Retourne le nombre de documents supprimés
    """

    documents = list_documents()

    if not documents or not pattern:
        return 0

    # Comptage des suppressions réussies
    success_count = 0

    # Filtrer les documents à supprimer
    docs_to_delete = []
    try:
        regex = re.compile(pattern)
        docs_to_delete = [doc for doc in documents if regex.search(doc["filename"])]
    except re.error as e:
        return 0

    # Supprimer les documents
    for doc in docs_to_delete:
        delete_document(doc["id"])
        # Petite pause pour éviter de surcharger l'API
        time.sleep(0.1)

    return success_count