import gradio as gr
import replicate
import os
from pytube import YouTube
from urllib.parse import urlparse, parse_qs
import requests
from PIL import Image
from io import BytesIO

# Clé API Replicate récupérée depuis Hugging Face secrets
REPLICATE_API_TOKEN = os.environ.get("REPLICATE_API_TOKEN")

# Nettoyage d'URL YouTube
def clean_youtube_url(url):
    parsed_url = urlparse(url)
    video_id = parse_qs(parsed_url.query).get("v", [None])[0]
    return f"https://www.youtube.com/watch?v={video_id}" if video_id else None

# Fonction principale de génération
def generate_from_youtube(url, style, ambiance):
    try:
        clean_url = clean_youtube_url(url)
        yt = YouTube(clean_url)
        title = yt.title
        prompt = f"Miniature YouTube pour un clip {style}, ambiance {ambiance}, titre : {title}"
    except Exception as e:
        prompt = f"Miniature YouTube {style}, ambiance {ambiance}"

    try:
        output = replicate.run(
            "stability-ai/stable-diffusion",
            input={
                "prompt": prompt,
                "image_dimensions": "512x512"
            }
        )

        # Débogage : afficher la sortie brute
        print("Résultat Replicate:", output)

        if not output or not isinstance(output, list) or not output[0].startswith("http"):
            return f"Erreur : le modèle n'a pas renvoyé une image valide.\nRetour : {output}"

        # Télécharger l'image et la convertir en objet PIL
        response = requests.get(output[0])
        image = Image.open(BytesIO(response.content))
        return image

    except Exception as e:
        return f"Erreur génération : {str(e)}"

# Choix de styles musicaux proposés
style_choices = ["Rap", "Afro", "Techno", "R&B", "Pop", "Reggae", "Drill", "Autre"]

# Interface Gradio
interface = gr.Interface(
    fn=generate_from_youtube,
    inputs=[
        gr.Textbox(label="Lien YouTube du clip"),
        gr.Dropdown(style_choices, label="Style musical"),
        gr.Textbox(label="Ambiance / Détail artistique (ex : sombre, festif, futuriste...)"),
    ],
    outputs=gr.Image(type="pil", label="Miniature générée"),
    title="YT Thumbnail Generator 🎬",
    description="Colle un lien YouTube et génère automatiquement une miniature stylisée."
)

# Lancement
interface.launch()