"""
generate.py — Script principal de génération vidéo (Texte → Vidéo)
Modèle utilisé : Wan 2.1 T2V 1.3B (Hugging Face Diffusers)

Usage :
    python generate.py --prompt "votre prompt" --output ../outputs/ma_video.mp4
"""

import argparse
import os
import sys


def _select_torch_dtype(torch_module):
    if not torch_module.cuda.is_available():
        return torch_module.float32

    if torch_module.cuda.is_bf16_supported():
        return torch_module.bfloat16

    return torch_module.float16

def generate_video(
    prompt: str,
    output_path: str,
    negative_prompt: str = "déformé, moche, flou, mauvaise qualité, artefacts, texte, filigrane",
    num_frames: int = 24,
    num_inference_steps: int = 25,
    height: int = 480,
    width: int = 832,
) -> str:
    """
    Génère une vidéo à partir d'un prompt texte et la sauvegarde dans le chemin spécifié.

    Args:
        prompt (str)             : Description textuelle de la vidéo à générer.
        output_path (str)        : Chemin de sauvegarde de la vidéo (ex: outputs/video.mp4).
        negative_prompt (str)    : Ce que l'on ne veut PAS voir dans la vidéo.
        num_frames (int)         : Nombre d'images à générer (24 ≈ 1 seconde à 24 fps).
        num_inference_steps (int): Nombre d'étapes de diffusion (plus = meilleure qualité, plus lent).
        height (int)             : Hauteur de la vidéo en pixels.
        width (int)              : Largeur de la vidéo en pixels.

    Returns:
        str: Le chemin de la vidéo générée, ou une chaîne vide en cas d'erreur.
    """
    try:
        import torch

        # Importation différée pour éviter les erreurs si les dépendances ne sont pas installées
        from diffusers import AutoencoderKLWan, WanPipeline
        from diffusers.utils import export_to_video

        print(f"[INFO] Chargement du modèle Wan 2.1 T2V 1.3B...")
        print(f"[INFO] (Le premier chargement peut prendre du temps — téléchargement des poids)")

        model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
        dtype = _select_torch_dtype(torch)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"[INFO] Appareil détecté : {device} | dtype : {dtype}")

        # Chargement du VAE (encodeur/décodeur vidéo)
        vae = AutoencoderKLWan.from_pretrained(
            model_id, subfolder="vae", torch_dtype=dtype
        )

        # Chargement du pipeline principal
        pipe = WanPipeline.from_pretrained(
            model_id, vae=vae, torch_dtype=dtype
        )

        if torch.cuda.is_available():
            # Optimisation mémoire : décharge les parties inactives du modèle vers la RAM
            pipe.enable_model_cpu_offload()
        else:
            pipe = pipe.to(device)

        print(f"[INFO] Modèle chargé. Génération de la vidéo...")
        print(f"[INFO] Prompt : {prompt}")

        # Génération des frames vidéo
        output = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_frames=num_frames,
            num_inference_steps=num_inference_steps,
        )

        # Création du dossier de sortie si nécessaire
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

        # Export en fichier vidéo MP4
        video_path = export_to_video(output.frames[0], output_path, fps=16)
        print(f"[SUCCESS] Vidéo sauvegardée : {video_path}")
        return video_path

    except ImportError as e:
        print(f"[ERREUR] Dépendance manquante : {e}")
        print("[AIDE] Installez les dépendances avec : pip install -r requirements.txt")
        return ""
    except OSError as e:
        print(f"[ERREUR] Problème d'environnement Python ou de bibliothèque native : {e}")
        print("[AIDE] Utilisez de préférence Python 3.10, 3.11 ou 3.12 avec une version compatible de PyTorch.")
        return ""
    except Exception as e:
        print(f"[ERREUR] Erreur lors de la génération : {e}")
        return ""


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Générateur de vidéo IA (Texte → Vidéo) avec Wan 2.1"
    )
    parser.add_argument(
        "--prompt", type=str, required=True,
        help="Description textuelle de la vidéo à générer"
    )
    parser.add_argument(
        "--output", type=str, default="../outputs/generated_video.mp4",
        help="Chemin de sortie pour la vidéo (défaut: ../outputs/generated_video.mp4)"
    )
    parser.add_argument(
        "--negative_prompt", type=str,
        default="déformé, moche, flou, mauvaise qualité, artefacts, texte, filigrane",
        help="Ce que l'on ne veut PAS voir dans la vidéo"
    )
    parser.add_argument(
        "--num_frames", type=int, default=24,
        help="Nombre d'images à générer (défaut: 24 ≈ 1.5 secondes)"
    )
    parser.add_argument(
        "--steps", type=int, default=25,
        help="Nombre d'étapes de diffusion (défaut: 25)"
    )

    args = parser.parse_args()

    result = generate_video(
        prompt=args.prompt,
        output_path=args.output,
        negative_prompt=args.negative_prompt,
        num_frames=args.num_frames,
        num_inference_steps=args.steps,
    )

    if result:
        print(f"\n✅ Génération réussie ! Vidéo disponible ici : {result}")
    else:
        print("\n❌ La génération a échoué. Consultez les messages d'erreur ci-dessus.")
        sys.exit(1)