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

import os
import sys
import argparse
import json
import base64
from PIL import Image
import io

def parse_args():
    parser = argparse.ArgumentParser(description="Test de prétraitement d'images")
    parser.add_argument("image_path", help="Chemin vers l'image à analyser")
    parser.add_argument("--save", "-s", action="store_true", help="Sauvegarder les images traitées")
    parser.add_argument("--verbose", "-v", action="store_true", help="Mode verbeux")
    return parser.parse_args()

def encoder_image_base64(image_path: str) -> str:
    """
    Encode une image en base64, avec optimisation de la taille si nécessaire.
    Implémentation identique à celle de LlamaVision._encoder_image_base64
    """
    try:
        # Vérifier la taille de l'image et la réduire si trop grande
        with Image.open(image_path) as img:
            # Afficher les informations de l'image originale
            print(f"Image originale: {image_path}")
            print(f"Format: {img.format}, Mode: {img.mode}, Taille: {img.size}")
            
            # Si l'image est trop grande, la redimensionner
            max_dim = 800  # Dimension maximale
            width, height = img.size
            
            if width > max_dim or height > max_dim:
                # Calculer le ratio pour conserver les proportions
                ratio = min(max_dim / width, max_dim / height)
                new_width = int(width * ratio)
                new_height = int(height * ratio)
                
                # Redimensionner l'image
                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
                print(f"Image redimensionnée: {new_width}x{new_height}")
            
            # Convertir en RGB si nécessaire (pour les formats comme PNG avec canal alpha)
            original_mode = img.mode
            if img.mode in ("RGBA", "LA", "P"):
                # Créer un fond blanc et composer l'image dessus pour gérer la transparence
                background = Image.new("RGB", img.size, (255, 255, 255))
                if img.mode == "P":
                    img = img.convert("RGBA")
                background.paste(img, mask=img.split()[3] if img.mode == "RGBA" else None)
                img = background
                print(f"Mode converti: {original_mode} -> RGB (avec fond blanc)")
            elif img.mode != "RGB":
                img = img.convert("RGB")
                print(f"Mode converti: {original_mode} -> RGB")
            
            # Sauvegarder temporairement l'image redimensionnée
            buffer = io.BytesIO()
            img.save(buffer, format="JPEG", quality=85)
            buffer.seek(0)
            
            # Sauvegarder l'image traitée si demandé
            if args.save:
                output_path = f"processed_{os.path.basename(image_path)}.jpg"
                img.save(output_path, format="JPEG", quality=85)
                print(f"Image traitée sauvegardée: {output_path}")
            
            # Encoder en base64
            encoded = base64.b64encode(buffer.read()).decode("utf-8")
            print(f"Taille du base64: {len(encoded)} caractères")
            return encoded
        
    except Exception as e:
        print(f"Erreur lors de l'optimisation de l'image: {str(e)}")
        try:
            # Seconde tentative avec une approche plus simple
            print("Tentative de secours...")
            with Image.open(image_path) as img:
                # Convertir directement en RGB quelle que soit l'image
                img = img.convert("RGB")
                buffer = io.BytesIO()
                img.save(buffer, format="JPEG", quality=75)
                buffer.seek(0)
                
                if args.save:
                    output_path = f"fallback_{os.path.basename(image_path)}.jpg"
                    img.save(output_path, format="JPEG", quality=75)
                    print(f"Image de secours sauvegardée: {output_path}")
                
                encoded = base64.b64encode(buffer.read()).decode("utf-8")
                print(f"Taille du base64 (secours): {len(encoded)} caractères")
                return encoded
        except Exception as e2:
            print(f"Deuxième erreur lors de l'optimisation de l'image: {str(e2)}")
            # Dernier recours: encoder l'image originale sans optimisation
            print("Dernier recours: encodage sans optimisation...")
            with open(image_path, "rb") as image_file:
                encoded = base64.b64encode(image_file.read()).decode("utf-8")
                print(f"Taille du base64 (brut): {len(encoded)} caractères")
                return encoded

def test_image_processing(image_path, verbose=False):
    """
    Teste le prétraitement d'image utilisé par LlamaVision
    """
    if not os.path.exists(image_path):
        print(f"Erreur: L'image {image_path} n'existe pas")
        return
    
    # Analyser l'image avec Pillow
    try:
        with Image.open(image_path) as img:
            print("\n=== INFORMATIONS SUR L'IMAGE ===")
            print(f"Format: {img.format}")
            print(f"Mode: {img.mode}")
            print(f"Taille: {img.size}")
            print(f"Palette: {hasattr(img, 'palette')}")
            if hasattr(img, 'info'):
                print(f"Info supplémentaires: {img.info.keys()}")
    except Exception as e:
        print(f"Erreur lors de l'analyse de l'image: {e}")
        return
    
    # Encoder l'image
    print("\n=== TRAITEMENT DE L'IMAGE ===")
    encoded = encoder_image_base64(image_path)
    
    # Sauvegarder des métadonnées
    metadata = {
        "filename": os.path.basename(image_path),
        "path": image_path,
        "base64_length": len(encoded),
        "first_20_chars": encoded[:20],
        "timestamp": __import__('datetime').datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    
    output_file = f"image_info_{os.path.basename(image_path)}.json"
    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(metadata, f, ensure_ascii=False, indent=2)
    
    print(f"\nMétadonnées enregistrées dans {output_file}")
    
    # Sauvegarder un extrait du base64
    if verbose:
        print("\n=== EXTRAIT DU BASE64 ===")
        print(encoded[:100] + "...")

if __name__ == "__main__":
    args = parse_args()
    test_image_processing(args.image_path, args.verbose)