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

"""
Script pour l'injection automatique de fichiers dans Private GPT.
Ce script scanne un répertoire source et injecte les nouveaux fichiers via l'API de Private GPT.
"""

import os
import sys
import time
import json
import argparse
import logging
import requests
from pathlib import Path
from typing import List, Set, Dict, Any
from datetime import datetime

# Configuration du logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('auto_ingest.log')
    ]
)
logger = logging.getLogger("auto_ingest")

# Extensions de fichiers supportées par défaut
DEFAULT_SUPPORTED_EXTENSIONS = {
    '.txt', '.pdf', '.csv', '.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx',
    '.md', '.html', '.htm', '.json', '.xml', '.rtf', '.odt', '.ods', '.odp'
}

class PrivateGPTIngestor:
    """Classe pour gérer l'ingestion de fichiers dans Private GPT."""

    def __init__(self, api_url: str = "http://localhost:8001",
                 processed_file: str = "processed_files.json"):
        """
        Initialise l'ingesteur.

        Args:
            api_url: URL de l'API Private GPT
            processed_file: Fichier pour stocker les fichiers déjà traités
        """
        self.api_url = api_url
        self.processed_file = processed_file
        self.processed_files = self._load_processed_files()

    def _load_processed_files(self) -> Set[str]:
        """Charge la liste des fichiers déjà traités."""
        try:
            if os.path.exists(self.processed_file):
                with open(self.processed_file, 'r', encoding='utf-8') as f:
                    return set(json.load(f))
            return set()
        except Exception as e:
            logger.error(f"Erreur lors du chargement des fichiers traités: {e}")
            return set()

    def _save_processed_files(self) -> None:
        """Sauvegarde la liste des fichiers déjà traités."""
        try:
            with open(self.processed_file, 'w', encoding='utf-8') as f:
                json.dump(list(self.processed_files), f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"Erreur lors de la sauvegarde des fichiers traités: {e}")

    def scan_directory(self, directory: str, extensions: Set[str] = None,
                      recursive: bool = True) -> List[str]:
        """
        Scanne un répertoire pour trouver des fichiers à injecter.

        Args:
            directory: Le répertoire à scanner
            extensions: Extensions de fichiers à prendre en compte
            recursive: Si True, scanne les sous-répertoires

        Returns:
            Liste des chemins des fichiers à injecter
        """
        if extensions is None:
            extensions = DEFAULT_SUPPORTED_EXTENSIONS

        files_to_ingest = []
        directory_path = Path(directory)

        if not directory_path.exists():
            logger.error(f"Le répertoire {directory} n'existe pas")
            return []

        logger.info(f"Scan du répertoire {directory}")

        # Fonction de scan
        def scan_dir(path: Path):
            for item in path.iterdir():
                if item.is_file() and any(item.name.lower().endswith(ext) for ext in extensions):
                    abs_path = str(item.absolute())
                    if abs_path not in self.processed_files:
                        files_to_ingest.append(abs_path)
                elif item.is_dir() and recursive:
                    scan_dir(item)

        scan_dir(directory_path)
        logger.info(f"Trouvé {len(files_to_ingest)} fichiers à injecter")
        return files_to_ingest

    def ingest_file(self, file_path: str) -> bool:
        """
        Injecte un fichier dans Private GPT via l'API.

        Args:
            file_path: Chemin du fichier à injecter

        Returns:
            True si l'injection a réussi, False sinon
        """
        logger.info(f"Injection du fichier: {file_path}")

        try:
            with open(file_path, 'rb') as f:
                files = {'file': (os.path.basename(file_path), f)}
                response = requests.post(f"{self.api_url}/v1/ingest/file", files=files)

                if response.status_code == 200:
                    logger.info(f"Injection réussie pour {file_path}")
                    self.processed_files.add(file_path)
                    self._save_processed_files()
                    return True
                else:
                    logger.error(f"Échec de l'injection pour {file_path}: {response.status_code} - {response.text}")
                    return False
        except Exception as e:
            logger.error(f"Erreur lors de l'injection de {file_path}: {e}")
            return False

    def list_documents(self) -> List[Dict[str, Any]]:
        """
        Liste les documents déjà injectés dans Private GPT.

        Returns:
            Liste des documents injectés
        """
        try:
            response = requests.get(f"{self.api_url}/v1/ingest/list")
            if response.status_code == 200:
                return response.json()
            else:
                logger.error(f"Échec de la récupération des documents: {response.status_code} - {response.text}")
                return []
        except Exception as e:
            logger.error(f"Erreur lors de la récupération des documents: {e}")
            return []

    def run_ingestion(self, directory: str, extensions: Set[str] = None,
                    recursive: bool = True, batch_size: int = 5,
                    delay: float = 2.0) -> None:
        """
        Exécute l'ingestion des fichiers d'un répertoire.

        Args:
            directory: Répertoire à scanner
            extensions: Extensions à prendre en compte
            recursive: Si True, scanne les sous-répertoires
            batch_size: Nombre de fichiers à injecter par lot
            delay: Délai entre chaque lot (en secondes)
        """
        files_to_ingest = self.scan_directory(directory, extensions, recursive)

        if not files_to_ingest:
            logger.info("Aucun nouveau fichier à injecter")
            return

        total_files = len(files_to_ingest)
        successful = 0
        failed = 0

        for i, file_path in enumerate(files_to_ingest):
            logger.info(f"Progression: {i+1}/{total_files}")

            if self.ingest_file(file_path):
                successful += 1
            else:
                failed += 1

            # Pause après chaque lot
            if (i + 1) % batch_size == 0 and i < total_files - 1:
                logger.info(f"Pause de {delay} secondes après le lot de {batch_size} fichiers")
                time.sleep(delay)

        logger.info(f"Ingestion terminée: {successful} réussis, {failed} échoués sur {total_files} fichiers")

def parse_args():
    """Parse les arguments de ligne de commande."""
    parser = argparse.ArgumentParser(description="Outil d'injection automatique pour Private GPT")

    parser.add_argument("--directory", "-d", type=str, required=True,
                        help="Répertoire contenant les fichiers à injecter")
    parser.add_argument("--api-url", type=str, default="http://localhost:8001",
                        help="URL de l'API Private GPT (défaut: http://localhost:8001)")
    parser.add_argument("--recursive", "-r", action="store_true", default=True,
                        help="Scanner récursivement les sous-répertoires")
    parser.add_argument("--extensions", "-e", type=str, nargs="+",
                        help="Extensions de fichiers à prendre en compte (ex: .pdf .txt)")
    parser.add_argument("--batch-size", "-b", type=int, default=5,
                        help="Nombre de fichiers à injecter par lot (défaut: 5)")
    parser.add_argument("--delay", type=float, default=2.0,
                        help="Délai entre les lots en secondes (défaut: 2.0)")
    parser.add_argument("--list", "-l", action="store_true",
                        help="Lister les documents déjà injectés et quitter")
    parser.add_argument("--watch", "-w", action="store_true",
                        help="Mode surveillance: vérifier périodiquement les nouveaux fichiers")
    parser.add_argument("--watch-interval", type=int, default=300,
                        help="Intervalle de surveillance en secondes (défaut: 300)")

    return parser.parse_args()

def main():
    """Fonction principale."""
    args = parse_args()

    ingestor = PrivateGPTIngestor(api_url=args.api_url)

    # Option pour lister les documents
    if args.list:
        documents = ingestor.list_documents()
        print(f"Documents déjà injectés ({len(documents)}):")
        for doc in documents:
            print(f"- {doc.get('doc_id')}: {doc.get('doc_metadata', {}).get('file_name', 'Inconnu')}")
        return

    # Conversion des extensions
    extensions = None
    if args.extensions:
        extensions = set()
        for ext in args.extensions:
            if not ext.startswith('.'):
                ext = '.' + ext
            extensions.add(ext.lower())

    # Mode surveillance
    if args.watch:
        logger.info(f"Mode surveillance activé. Vérification toutes les {args.watch_interval} secondes")
        try:
            while True:
                start_time = datetime.now()
                logger.info(f"Démarrage d'un scan à {start_time.strftime('%H:%M:%S')}")

                ingestor.run_ingestion(
                    directory=args.directory,
                    extensions=extensions,
                    recursive=args.recursive,
                    batch_size=args.batch_size,
                    delay=args.delay
                )

                # Calcul du temps à attendre
                elapsed = (datetime.now() - start_time).total_seconds()
                wait_time = max(0, args.watch_interval - elapsed)

                if wait_time > 0:
                    logger.info(f"En attente pendant {wait_time:.1f} secondes jusqu'au prochain scan")
                    time.sleep(wait_time)
        except KeyboardInterrupt:
            logger.info("Arrêt du mode surveillance")
    else:
        # Exécution unique
        ingestor.run_ingestion(
            directory=args.directory,
            extensions=extensions,
            recursive=args.recursive,
            batch_size=args.batch_size,
            delay=args.delay
        )

if __name__ == "__main__":
    main()
