import streamlit as st
from PIL import Image
import torch
import tempfile
import os
import io
import base64

# Configuration de la page - IMPORTANT pour Hugging Face Spaces
st.set_page_config(
    page_title="Video AI Generator",
    page_icon="🎥",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Configuration spéciale pour Hugging Face Spaces
if 'HF_SPACE' in os.environ:
    # Désactiver les vérifications CSRF pour HF Spaces
    st.set_option('server.enableCORS', True)
    st.set_option('server.enableXsrfProtection', False)

# Configuration pour éviter les erreurs de mémoire
@st.cache_resource
def load_pipeline():
    """Charge le pipeline avec optimisations mémoire"""
    try:
        from diffusers import StableDiffusionImg2ImgPipeline
        
        # Utiliser un modèle plus léger ou optimisé
        model_id = "runwayml/stable-diffusion-v1-5"
        
        pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            use_safetensors=True,
        )
        
        # Optimisations mémoire
        if torch.cuda.is_available():
            pipeline = pipeline.to("cuda")
            pipeline.enable_memory_efficient_attention()
            pipeline.enable_attention_slicing()
        else:
            pipeline = pipeline.to("cpu")
            
        return pipeline
    except Exception as e:
        st.error(f"Erreur lors du chargement du modèle : {e}")
        return None

def safe_file_upload():
    """Upload de fichier sécurisé avec gestion d'erreurs"""
    try:
        # Upload avec paramètres optimisés pour HF Spaces
        uploaded_file = st.file_uploader(
            "Choisissez une image",
            type=['png', 'jpg', 'jpeg'],
            accept_multiple_files=False,
            key="image_uploader",
            help="Formats supportés: PNG, JPG, JPEG (max 10MB recommandé)"
        )
        
        if uploaded_file is not None:
            # Vérifier la taille du fichier
            file_size = len(uploaded_file.getvalue())
            if file_size > 10 * 1024 * 1024:  # 10MB
                st.warning(f"⚠️ Fichier volumineux ({file_size/1024/1024:.1f}MB). Réduisez la taille pour de meilleures performances.")
            
            # Charger l'image avec gestion d'erreur
            try:
                image = Image.open(uploaded_file)
                st.success(f"✅ Image chargée: {image.size[0]}x{image.size[1]} pixels")
                return image, uploaded_file
            except Exception as img_error:
                st.error(f"❌ Erreur lors du chargement de l'image: {img_error}")
                return None, None
                
    except Exception as upload_error:
        st.error(f"❌ Erreur d'upload: {upload_error}")
        st.info("💡 Essayez de rafraîchir la page ou utilisez un fichier plus petit")
        return None, None
    
    return None, None

def process_image(image, prompt, strength=0.75, guidance_scale=7.5):
    """Traite l'image avec le pipeline"""
    try:
        pipeline = load_pipeline()
        if pipeline is None:
            return None
            
        # Redimensionner l'image pour éviter les problèmes de mémoire
        max_size = 512
        if image.width > max_size or image.height > max_size:
            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        
        # S'assurer que l'image est en RGB
        if image.mode != "RGB":
            image = image.convert("RGB")
        
        # Générer l'image
        with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
            result = pipeline(
                prompt=prompt,
                image=image,
                strength=strength,
                guidance_scale=guidance_scale,
                num_inference_steps=20,  # Réduire pour plus de vitesse
                height=image.height,
                width=image.width
            )
        
        return result.images[0]
    except Exception as e:
        st.error(f"Erreur lors de la génération : {e}")
        return None

# Interface utilisateur
def main():
    st.title("🎥 Générateur Video AI")
    st.markdown("Transformez vos images avec l'IA")
    
    # Afficher le statut de l'environnement
    if 'HF_SPACE' in os.environ:
        st.info("🚀 Application déployée sur Hugging Face Spaces")
    
    # Sidebar pour les paramètres
    with st.sidebar:
        st.header("⚙️ Paramètres")
        
        strength = st.slider(
            "Force de transformation",
            min_value=0.1,
            max_value=1.0,
            value=0.75,
            step=0.05,
            help="Plus élevé = plus de changements"
        )
        
        guidance_scale = st.slider(
            "Échelle de guidance",
            min_value=1.0,
            max_value=20.0,
            value=7.5,
            step=0.5,
            help="Plus élevé = plus fidèle au prompt"
        )
        
        st.markdown("---")
        st.markdown("💡 **Conseils:**")
        st.markdown("- Utilisez des images < 10MB")
        st.markdown("- Format recommandé: PNG, JPG")
        st.markdown("- Résolution max: 512x512")
    
    # Section principale
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.header("📤 Image d'entrée")
        
        # Upload de fichier sécurisé
        input_image, uploaded_file = safe_file_upload()
        
        if input_image is not None:
            # Redimensionner pour l'affichage
            display_image = input_image.copy()
            if display_image.width > 400:
                display_image.thumbnail((400, 400), Image.Resampling.LANCZOS)
            
            st.image(display_image, caption="Image d'origine", use_column_width=True)
            
            # Zone de texte pour le prompt
            prompt = st.text_area(
                "🎨 Décrivez la transformation souhaitée:",
                placeholder="Ex: transformer en anime style, ajouter des couleurs vives, style cyberpunk...",
                height=100,
                key="prompt_input"
            )
            
            # Bouton de génération
            if st.button("🚀 Générer", type="primary", use_container_width=True):
                if prompt.strip():
                    with st.spinner("Génération en cours... Cela peut prendre quelques minutes."):
                        result_image = process_image(
                            input_image, 
                            prompt, 
                            strength, 
                            guidance_scale
                        )
                    
                    if result_image:
                        # Stocker le résultat dans la session
                        st.session_state['result_image'] = result_image
                        st.success("✅ Génération terminée!")
                        st.rerun()
                    else:
                        st.error("❌ Erreur lors de la génération")
                else:
                    st.warning("⚠️ Veuillez entrer une description")
    
    with col2:
        st.header("🎨 Résultat")
        
        # Afficher le résultat s'il existe
        if 'result_image' in st.session_state:
            result_image = st.session_state['result_image']
            
            # Redimensionner pour l'affichage
            display_result = result_image.copy()
            if display_result.width > 400:
                display_result.thumbnail((400, 400), Image.Resampling.LANCZOS)
            
            st.image(display_result, caption="Image transformée", use_column_width=True)
            
            # Bouton de téléchargement
            try:
                buffered = io.BytesIO()
                result_image.save(buffered, format="PNG")
                
                st.download_button(
                    label="💾 Télécharger le résultat",
                    data=buffered.getvalue(),
                    file_name="result_ai.png",
                    mime="image/png",
                    use_container_width=True
                )
            except Exception as e:
                st.error(f"Erreur lors de la préparation du téléchargement: {e}")
        else:
            st.info("👈 Uploadez une image et entrez un prompt pour commencer")
    
    # Section de dépannage
    with st.expander("🔧 Problèmes d'upload?"):
        st.markdown("""
        **Si vous ne pouvez pas charger d'images:**
        
        1. **Rafraîchissez** la page (F5)
        2. **Réduisez** la taille de votre image (< 5MB)
        3. **Changez** le format (essayez JPG au lieu de PNG)
        4. **Utilisez** un autre navigateur
        5. **Attendez** quelques secondes et réessayez
        
        **Formats supportés:** PNG, JPG, JPEG
        **Taille recommandée:** < 10MB
        """)
    
    # Instructions
    with st.expander("📖 Instructions d'utilisation"):
        st.markdown("""
        ### Comment utiliser cette application:
        
        1. **Uploadez une image** dans la section de gauche
        2. **Décrivez la transformation** que vous souhaitez
        3. **Ajustez les paramètres** dans la sidebar si nécessaire
        4. **Cliquez sur Générer** et patientez
        5. **Téléchargez le résultat** une fois terminé
        
        ### Exemples de prompts:
        - "transformer en style anime"
        - "ajouter des couleurs arc-en-ciel"
        - "style cyberpunk néon"
        - "portrait artistique à l'aquarelle"
        - "paysage fantastique avec des dragons"
        """)

if __name__ == "__main__":
    main()