llm_ticket3/extract_single_message.py

# extract_single_message.py

import json
import os
import sys
import argparse
from bs4 import BeautifulSoup, Tag

def parse_args():
    """Parse les arguments de ligne de commande."""
    parser = argparse.ArgumentParser(description="Extrait un message spécifique et ses images depuis un fichier JSON.")
    
    parser.add_argument("--input", "-i", 
                        help="Chemin vers le fichier JSON contenant les messages (par défaut: messages_raw.json)",
                        default="messages_raw.json")
    
    parser.add_argument("--message-id", "-m", 
                        help="ID du message à extraire (par défaut: 228942)",
                        type=int, default=228942)
    
    parser.add_argument("--output", "-o", 
                        help="Répertoire de sortie (par défaut: extracted_message)",
                        default="extracted_message")
    
    parser.add_argument("--create-test", "-t", 
                        help="Créer un fichier de test si le fichier d'entrée n'existe pas",
                        action="store_true")
    
    return parser.parse_args()

def create_test_data(input_file, message_id):
    """Crée un fichier de test contenant un message avec une image."""
    print(f"⚠️ Le fichier {input_file} n'existe pas. Création d'un exemple de test...")
    
    test_data = {
        "messages": [
            {
                "id": message_id,
                "body": "<p>Voici une image de test:</p><img src='https://example.com/image.jpg' alt='Test'>",
                "author_id": [1, "Utilisateur Test"],
                "date": "2023-01-01 12:00:00",
                "subject": "Message de test"
            }
        ]
    }
    
    try:
        # Créer le répertoire parent si nécessaire
        parent_dir = os.path.dirname(input_file)
        if parent_dir and not os.path.exists(parent_dir):
            os.makedirs(parent_dir, exist_ok=True)
            
        with open(input_file, "w", encoding="utf-8") as f:
            json.dump(test_data, f, indent=2, ensure_ascii=False)
        
        print(f"✅ Fichier de test créé : {input_file}")
        return True
    except Exception as e:
        print(f"❌ Erreur lors de la création du fichier de test : {e}")
        return False

def extract_message(input_file, message_id, output_dir, create_test=False):
    """Extrait un message spécifique et ses images depuis un fichier JSON."""
    # Vérifier si le fichier d'entrée existe
    if not os.path.exists(input_file):
        if create_test:
            created = create_test_data(input_file, message_id)
            if not created:
                print(f"❌ Impossible de créer le fichier de test {input_file}")
                return False
        else:
            print(f"❌ Le fichier {input_file} n'existe pas. Utilisez --create-test pour créer un exemple.")
            return False

    try:
        # Créer le répertoire de sortie
        os.makedirs(output_dir, exist_ok=True)
        
        # Charger les données
        with open(input_file, "r", encoding="utf-8") as f:
            data = json.load(f)

        messages = data.get("messages", [])
        found = next((m for m in messages if m["id"] == message_id), None)

        if not found:
            print(f"❌ Message ID {message_id} non trouvé.")
            return False

        print(f"✅ Message ID {message_id} trouvé.")

        # 1. Sauvegarder le HTML brut
        body_html = found.get("body", "")
        html_path = os.path.join(output_dir, f"message_{message_id}.html")
        with open(html_path, "w", encoding="utf-8") as f:
            f.write(body_html)
        print(f"📄 HTML sauvegardé dans : {html_path}")

        # 2. Extraire les images depuis le HTML
        soup = BeautifulSoup(body_html, "html.parser")
        image_urls = []
        
        # Extraction sécurisée des URLs d'images
        for img in soup.find_all("img"):
            if isinstance(img, Tag) and img.has_attr("src"):
                image_urls.append(img["src"])
        
        image_data = {
            "message_id": message_id,
            "author": found.get("author_id", [None, None])[1],
            "date": found.get("date"),
            "subject": found.get("subject"),
            "image_urls": image_urls
        }

        # 3. Sauvegarder sous forme JSON
        json_path = os.path.join(output_dir, f"message_{message_id}_meta.json")
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(image_data, f, indent=2, ensure_ascii=False)
        print(f"🧾 Métadonnées sauvegardées dans : {json_path}")

        print("\n✅ Extraction terminée.")
        print("🖼️ Images détectées :")
        for url in image_urls:
            print(f" - {url}")
        
        return True
    except Exception as e:
        print(f"❌ Erreur lors de l'extraction : {e}")
        return False

if __name__ == "__main__":
    args = parse_args()
    success = extract_message(
        input_file=args.input,
        message_id=args.message_id,
        output_dir=args.output,
        create_test=args.create_test
    )
    
    if not success:
        sys.exit(1)