import os
import argparse
from pathlib import Path
from PIL import Image
import torch

# =====================================================================
# 🔥 TRUCO TÉCNICO: PARCHE ULTRA-AGRESIVO PARA FORZAR CPU 🔥
# =====================================================================
# Desactivamos por completo las llamadas de inicialización de CUDA
torch.cuda.is_available = lambda: False
torch.cuda.set_device = lambda *args, **kwargs: None
torch.cuda.init = lambda *args, **kwargs: None

# Forzamos a PyTorch a descargar y mapear los archivos directo a la CPU
original_load = torch.load
torch.load = lambda *args, **kwargs: original_load(*args, **kwargs, map_location=torch.device('cpu'))

# Si la librería interna intenta mover un Tensor a CUDA, lo mandamos a CPU
original_tensor_to = torch.Tensor.to
def patched_tensor_to(self, *args, **kwargs):
    new_args = tuple("cpu" if isinstance(a, str) and "cuda" in a else a for a in args)
    if "device" in kwargs and isinstance(kwargs["device"], str) and "cuda" in kwargs["device"]:
        kwargs["device"] = "cpu"
    return original_tensor_to(self, *new_args, **kwargs)
torch.Tensor.to = patched_tensor_to
torch.Tensor.cuda = lambda self, *args, **kwargs: self.cpu()

# Si la librería intenta mover todo el Modelo a CUDA, lo mandamos a CPU
original_module_to = torch.nn.Module.to
def patched_module_to(self, *args, **kwargs):
    new_args = tuple("cpu" if isinstance(a, str) and "cuda" in a else a for a in args)
    if "device" in kwargs and isinstance(kwargs["device"], str) and "cuda" in kwargs["device"]:
        kwargs["device"] = "cpu"
    return original_module_to(self, *new_args, **kwargs)
torch.nn.Module.to = patched_module_to
torch.nn.Module.cuda = lambda self, *args, **kwargs: self.cpu()
# =====================================================================

# Ahora que PyTorch está "blindado", importamos AuraSR de forma segura
from aura_sr import AuraSR

def main():
    # Configurar los argumentos por línea de comandos
    parser = argparse.ArgumentParser(
        description="Mejora por lotes de imágenes deportivas forzando el uso de CPU."
    )
    parser.add_argument(
        "-d", "--input_dir", 
        required=True, 
        help="Ruta del directorio de origen que contiene las imágenes originales."
    )
    parser.add_argument(
        "-dd", "--output_dir", 
        required=True, 
        help="Ruta del directorio de destino donde se guardarán las fotos finales mejoradas."
    )
    
    args = parser.parse_args()
    
    input_path = Path(args.input_dir)
    output_path = Path(args.output_dir)
    
    # Validar directorios
    if not input_path.exists() or not input_path.is_dir():
        print(f"\n[ERROR] El directorio de origen '{input_path}' no existe o no es válido.")
        return
        
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Extensiones de imagen soportadas
    valid_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tiff'}
    image_files = [f for f in input_path.iterdir() if f.is_file() and f.suffix.lower() in valid_extensions]
    
    if not image_files:
        print(f"\n[INFO] No se encontraron imágenes válidas en '{input_path}'.")
        return
        
    print(f"\n Se encontraron {len(image_files)} imágenes para procesar.")
    print("Forzando entorno: CPU (Simulación activa)")
    print("Cargando el modelo AuraSR v2 en la memoria RAM...")
    
    try:
        # Carga normal (el parche se encarga de que no busque CUDA internamente)
        aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
    except Exception as e:
        print(f"\n[ERROR] No se pudo cargar el modelo AuraSR: {e}")
        return

    print("\n¡Comenzando la mejora por lotes en CPU!")
    print("-" * 70)
    
    for idx, img_file in enumerate(image_files, 1):
        print(f"[{idx}/{len(image_files)}] Procesando: {img_file.name}...", end="", flush=True)
        try:
            with Image.open(img_file) as img:
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # Ejecutar escalado 4x
                img_enhanced = aura_sr.upscale_4x_overlapped(img)
                
                # Guardar resultado
                out_file = output_path / f"{img_file.stem}_hd.png"
                img_enhanced.save(out_file, "PNG")
                
                print(f" -> ¡Hecho! Guardado como: {out_file.name}")
        except Exception as e:
            print(f" -> [ERROR] No se pudo procesar: {e}")

    print("-" * 70)
    print(f" Proceso completado. Revisa tus fotos en: '{output_path.resolve()}'\n")

if __name__ == "__main__":
    main()