import os
import sys
import glob
from PIL import Image
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer

def setup_upscaler():
    """Configure l'upscaler Real-ESRGAN."""
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
    model_path = 'RealESRGAN_x4plus.pth'
    
    # Téléchargement automatique du modèle si absent (géré par RealESRGANer ou wget)
    if not os.path.exists(model_path):
        import subprocess
        print("📥 Téléchargement du modèle Real-ESRGAN...")
        subprocess.run(["wget", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"])
        
    upsampler = RealESRGANer(
        scale=4,
        model_path=model_path,
        model=model,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=True # Utilise FP16 pour les GPUs modernes (A10, A100, etc.)
    )
    return upsampler

def main():
    input_dir = "input_images"
    output_dir = "output_images"
    os.makedirs(output_dir, exist_ok=True)
    
    print("🚀 Initialisation de l'upscaler AI...")
    upsampler = setup_upscaler()
    
    images = glob.glob(os.path.join(input_dir, "*"))
    print(f"📂 {len(images)} images à traiter.")
    
    for img_path in images:
        name = os.path.basename(img_path)
        print(f"🪄 Processing {name}...")
        try:
            import cv2
            img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
            output, _ = upsampler.enhance(img, outscale=4)
            cv2.imwrite(os.path.join(output_dir, name), output)
            print(f"✅ Terminé : {name}")
        except Exception as e:
            print(f"❌ Erreur sur {name} : {e}")

if __name__ == "__main__":
    # Installation des dépendances si nécessaire (sur l'instance Lambda)
    import subprocess
    print("📦 Installation des dépendances AI...")
    subprocess.run(["pip", "install", "basicsr", "realesrgan", "opencv-python"])
    
    main()