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

"""
Script pour tester directement l'agent de tri d'images (AgentImageSorter) avec Pixtral-Large.
Permet d'évaluer rapidement si une image est pertinente pour un ticket.
"""

import os
import sys
import argparse
import glob
from typing import List

from llm_classes.pixtral_large import PixtralLarge
from agents.pixtral_large.agent_image_sorter import AgentImageSorter
from utils.image_dedup import filtrer_images_uniques

def get_images_in_extraction(extraction_path: str) -> List[str]:
    """
    Récupère toutes les images dans un répertoire d'extraction
    
    Args:
        extraction_path: Chemin vers le répertoire d'extraction
        
    Returns:
        Liste des chemins d'accès aux images
    """
    # Extensions d'images courantes
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.gif', '*.webp']
    
    # Liste pour stocker les chemins d'images
    images = []
    
    # Chercher dans le répertoire principal
    for extension in image_extensions:
        pattern = os.path.join(extraction_path, extension)
        images.extend(glob.glob(pattern))
    
    # Chercher dans le sous-répertoire "attachments"
    attachments_path = os.path.join(extraction_path, "attachments")
    if os.path.exists(attachments_path) and os.path.isdir(attachments_path):
        for extension in image_extensions:
            pattern = os.path.join(attachments_path, extension)
            images.extend(glob.glob(pattern))
        
        # Chercher dans les sous-répertoires de "attachments" (si des ID sont utilisés)
        for subdir in os.listdir(attachments_path):
            subdir_path = os.path.join(attachments_path, subdir)
            if os.path.isdir(subdir_path):
                for extension in image_extensions:
                    pattern = os.path.join(subdir_path, extension)
                    images.extend(glob.glob(pattern))
    
    return images

def main():
    # Configuration de l'analyseur d'arguments
    parser = argparse.ArgumentParser(description="Tester l'agent de tri d'images directement.")
    parser.add_argument("ticket_id", help="ID du ticket à analyser (ex: T1234)")
    parser.add_argument("--image", help="Chemin spécifique vers une image à analyser")
    parser.add_argument("--output_dir", default="output", help="Répertoire de sortie contenant les tickets")
    parser.add_argument("--no-dedup", action="store_true", help="Désactiver le préfiltrage des doublons")
    parser.add_argument("--seuil", type=int, default=5, help="Seuil de similarité pour détecter les doublons (0-10, défaut=5)")
    
    # Analyser les arguments
    args = parser.parse_args()
    
    # Construire le chemin vers le ticket
    ticket_path = os.path.join(args.output_dir, f"ticket_{args.ticket_id}")
    
    # Vérifier que le répertoire du ticket existe
    if not os.path.exists(ticket_path):
        print(f"ERREUR: Le ticket {args.ticket_id} n'existe pas dans {args.output_dir}")
        return 1
    
    # Rechercher la dernière extraction (la plus récente)
    extractions = [d for d in os.listdir(ticket_path) if os.path.isdir(os.path.join(ticket_path, d)) and d.startswith(args.ticket_id)]
    if not extractions:
        print(f"ERREUR: Aucune extraction trouvée pour le ticket {args.ticket_id}")
        return 1
    
    # Trier par ordre décroissant pour avoir la plus récente en premier
    extractions.sort(reverse=True)
    latest_extraction = extractions[0]
    extraction_path = os.path.join(ticket_path, latest_extraction)
    
    print(f"Utilisation de l'extraction: {latest_extraction}")
    
    # Initialiser le modèle Pixtral-Large
    try:
        print("Initialisation du modèle Pixtral-Large...")
        model = PixtralLarge()
        agent = AgentImageSorter(model)
        print("Agent de tri d'images initialisé avec succès")
    except Exception as e:
        print(f"ERREUR: Impossible d'initialiser le modèle: {str(e)}")
        return 1
    
    # Si une image spécifique est fournie
    if args.image:
        image_path = args.image
        if not os.path.exists(image_path):
            print(f"ERREUR: L'image spécifiée n'existe pas: {image_path}")
            return 1
        
        print(f"Analyse de l'image: {os.path.basename(image_path)}")
        resultat = agent.executer(image_path)
        
        # Afficher le résultat
        print("\nRésultat de l'analyse:")
        print(f"Pertinence: {'OUI' if resultat.get('is_relevant', False) else 'NON'}")
        print(f"Raison: {resultat.get('reason', 'Non spécifiée')}")
        
    else:
        # Récupérer toutes les images de l'extraction
        images = get_images_in_extraction(extraction_path)
        
        if not images:
            print(f"Aucune image trouvée dans l'extraction {latest_extraction}")
            return 1
        
        print(f"Nombre d'images trouvées: {len(images)}")
        
        # Appliquer le préfiltrage de doublons si activé
        if not args.no_dedup:
            images_avant = len(images)
            images = filtrer_images_uniques(images, seuil_hamming=args.seuil, ticket_id=args.ticket_id)
            images_apres = len(images)
            
            if images_avant > images_apres:
                print(f"Préfiltrage des doublons: {images_avant} → {images_apres} images ({images_avant - images_apres} doublons supprimés)")
            else:
                print("Préfiltrage terminé: aucun doublon détecté")
        
        # Analyser chaque image
        results = []
        for image_path in images:
            print(f"\nAnalyse de l'image: {os.path.basename(image_path)}")
            resultat = agent.executer(image_path)
            results.append(resultat)
            
            # Afficher le résultat
            print(f"Pertinence: {'OUI' if resultat.get('is_relevant', False) else 'NON'}")
            print(f"Raison: {resultat.get('reason', 'Non spécifiée')}")
        
        # Afficher un résumé à la fin
        pertinentes = sum(1 for r in results if r.get('is_relevant', False))
        print(f"\nRésumé: {pertinentes}/{len(results)} images pertinentes")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())