import gradio as gr
import os
from PIL import Image, ImageEnhance
import requests
import io
import gc
import json
from typing import Tuple, Optional, Dict, Any
import logging
from dotenv import load_dotenv

# Configuration du logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Chargement des variables d'environnement
load_dotenv()

# Styles artistiques complets
ART_STYLES = {
    # Styles Réalistes
    "Ultra Réaliste": {
        "prompt_prefix": "ultra realistic photograph, stunning photorealistic quality, unreal engine 5 quality, octane render, ray tracing, volumetric lighting, subsurface scattering, 8k UHD, cinema quality, masterpiece, perfect composition, award winning photography",
        "negative_prompt": "artificial, digital art, illustration, painting, drawing, artistic, cartoon, anime, unreal, fake, low quality, blurry, soft, deformed"
    },
    "Photoréaliste": {
        "prompt_prefix": "hyperrealistic photograph, extremely detailed, studio quality, professional photography, 8k uhd",
        "negative_prompt": "artistic, painterly, abstract, cartoon, illustration, low quality"
    },
    "Expressionniste": {
        "prompt_prefix": "expressive painting style, intense emotional art, bold brushstrokes, vibrant colors, van gogh inspired",
        "negative_prompt": "realistic, subtle, photographic, clean lines, digital art"
    },
    "Impressionniste": {
        "prompt_prefix": "impressionist painting style, soft light, visible brushstrokes, outdoor scene, monet inspired",
        "negative_prompt": "sharp details, high contrast, digital, modern"
    },
    "Art Abstrait": {
        "prompt_prefix": "abstract art, geometric shapes, non-representational, kandinsky style, pure artistic expression",
        "negative_prompt": "realistic, figurative, photographic, literal"
    },
    "Art Moderne": {
        "prompt_prefix": "modern art style poster, professional design, contemporary aesthetic",
        "negative_prompt": "traditional, cluttered, busy design, vintage"
    },
    "Minimaliste": {
        "prompt_prefix": "minimalist design poster, clean composition, elegant simplicity",
        "negative_prompt": "complex, detailed, ornate, busy, cluttered"
    }
}

# Paramètres de composition
COMPOSITION_PARAMS = {
    "Layouts": {
        "Centré": "centered composition, balanced layout, harmonious arrangement",
        "Asymétrique": "dynamic asymmetrical composition, creative balance",
        "Grille": "grid-based layout, structured composition, organized design",
        "Diagonal": "diagonal dynamic composition, energetic flow",
        "Minimaliste": "minimal composition, lots of whitespace, elegant spacing"
    },
    "Ambiances": {
        "Dramatique": "dramatic lighting, high contrast, intense mood",
        "Doux": "soft lighting, gentle atmosphere, subtle mood",
        "Vibrant": "vibrant colors, energetic mood, dynamic atmosphere",
        "Mystérieux": "mysterious atmosphere, moody lighting, enigmatic feel",
        "Serein": "peaceful atmosphere, calm mood, tranquil setting"
    },
    "Palette": {
        "Monochrome": "monochromatic color scheme, sophisticated tones",
        "Contrasté": "high contrast color palette, bold color combinations",
        "Pastel": "soft pastel color palette, gentle colors",
        "Terre": "earthy color palette, natural tones",
        "Néon": "neon color palette, vibrant glowing colors"
    }
}

class ImageGenerator:
    def __init__(self):
        self.API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
        token = os.getenv('HUGGINGFACE_TOKEN')
        if not token:
            logger.error("HUGGINGFACE_TOKEN non trouvé!")
        self.headers = {"Authorization": f"Bearer {token}"}
        logger.info("ImageGenerator initialisé")

    def _build_prompt(self, params: Dict[str, Any]) -> str:
        try:
            style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
            
            # Construction du prompt de base
            base_prompt = f"{params['subject']}"
            if params.get('title'):
                base_prompt += f", with text '{params['title']}'"

            # Ajout des éléments de composition
            composition_elements = [
                style_info['prompt_prefix'],
                COMPOSITION_PARAMS['Layouts'][params['layout']],
                COMPOSITION_PARAMS['Ambiances'][params['ambiance']],
                COMPOSITION_PARAMS['Palette'][params['palette']]
            ]

            enhanced_prompt = f"{base_prompt}, {', '.join(composition_elements)}"
            
            return enhanced_prompt

        except Exception as e:
            logger.error(f"Erreur dans la construction du prompt: {str(e)}")
            return f"{style_info['prompt_prefix']}, {params['subject']}"

    def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
        try:
            logger.info(f"Début de génération avec paramètres: {json.dumps(params, indent=2)}")

            if 'Bearer None' in self.headers['Authorization']:
                return None, "⚠️ Erreur: Token Hugging Face non configuré"

            prompt = self._build_prompt(params)
            style_info = ART_STYLES[params["style"]]

            payload = {
                "inputs": prompt,
                "parameters": {
                    "negative_prompt": style_info["negative_prompt"],
                    "num_inference_steps": min(int(35 * (params["quality"]/100)), 40),
                    "guidance_scale": min(7.5 * (params["creativity"]/10), 10.0),
                    "width": 1024 if params.get("quality", 35) > 40 else 768,
                    "height": 1024 if params["orientation"] == "Portrait" else 768
                }
            }

            logger.debug(f"Payload: {json.dumps(payload, indent=2)}")

            response = requests.post(
                self.API_URL,
                headers=self.headers,
                json=payload,
                timeout=45
            )

            if response.status_code == 200:
                image = Image.open(io.BytesIO(response.content))
                # Post-traitement basé sur les paramètres
                image = self._enhance_image(image, params)
                return image, "✨ Création réussie!"
            else:
                error_msg = f"⚠️ Erreur API {response.status_code}: {response.text}"
                logger.error(error_msg)
                return None, error_msg

        except Exception as e:
            error_msg = f"⚠️ Erreur: {str(e)}"
            logger.exception("Erreur pendant la génération:")
            return None, error_msg
        finally:
            gc.collect()

    def _enhance_image(self, image: Image.Image, params: Dict[str, Any]) -> Image.Image:
        """Applique des améliorations post-génération à l'image"""
        try:
            # Ajustement du contraste
            if params.get("contrast", 5) != 5:
                enhancer = ImageEnhance.Contrast(image)
                factor = params["contrast"] / 5
                image = enhancer.enhance(factor)

            # Ajustement de la saturation
            if params.get("saturation", 5) != 5:
                enhancer = ImageEnhance.Color(image)
                factor = params["saturation"] / 5
                image = enhancer.enhance(factor)

            return image
        except Exception as e:
            logger.warning(f"Erreur lors de l'amélioration de l'image: {e}")
            return image

def create_interface():
    logger.info("Création de l'interface Gradio")

    css = """
    .container { max-width: 1200px; margin: auto; }
    .welcome { text-align: center; margin: 20px 0; padding: 20px; background: #1e293b; border-radius: 10px; }
    .controls-group { background: #2d3748; padding: 15px; border-radius: 5px; margin: 10px 0; }
    .advanced-controls { background: #374151; padding: 12px; border-radius: 5px; margin: 8px 0; }
    """

    generator = ImageGenerator()

    with gr.Blocks(css=css) as app:
        gr.HTML("""
            <div class="welcome">
                <h1>🎨 Equity Artisan 3.0</h1>
                <p>Assistant de création d'affiches professionnelles</p>
            </div>
        """)

        with gr.Column(elem_classes="container"):
            # Guide d'utilisation
            gr.Markdown("""
                ### 🎯 Guide d'utilisation
                1. Choisissez le format et l'orientation de votre affiche
                2. Sélectionnez un style artistique et une composition
                3. Décrivez votre vision dans "Description"
                4. Ajustez les paramètres fins selon vos besoins
                5. Cliquez sur "Générer" !
            """)

            # Format et Orientation
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 📐 Format et Orientation")
                with gr.Row():
                    format_size = gr.Dropdown(
                        choices=["A4", "A3", "A2", "A1", "A0"],
                        value="A4",
                        label="Format"
                    )
                    orientation = gr.Radio(
                        choices=["Portrait", "Paysage"],
                        value="Portrait",
                        label="Orientation"
                    )

            # Style et Composition
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 🎨 Style et Composition")
                with gr.Row():
                    style = gr.Dropdown(
                        choices=list(ART_STYLES.keys()),
                        value="Art Moderne",
                        label="Style artistique"
                    )
                    layout = gr.Dropdown(
                        choices=list(COMPOSITION_PARAMS["Layouts"].keys()),
                        value="Centré",
                        label="Composition"
                    )

                with gr.Row():
                    ambiance = gr.Dropdown(
                        choices=list(COMPOSITION_PARAMS["Ambiances"].keys()),
                        value="Dramatique",
                        label="Ambiance"
                    )
                    palette = gr.Dropdown(
                        choices=list(COMPOSITION_PARAMS["Palette"].keys()),
                        value="Contrasté",
                        label="Palette"
                    )

            # Contenu
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 📝 Contenu")
                subject = gr.Textbox(
                    label="Description",
                    placeholder="Ex: Une affiche moderne pour un festival de musique...",
                    lines=3
                )
                title = gr.Textbox(
                    label="Titre (optionnel)",
                    placeholder="Le titre qui apparaîtra sur l'affiche..."
                )

            # Contrôles avancés
            with gr.Group(elem_classes="advanced-controls"):
                gr.Markdown("### 🎯 Paramètres Avancés")
                with gr.Row():
                    detail_level = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=7,
                        step=1,
                        label="Niveau de Détail"
                    )
                    contrast = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=5,
                        step=1,
                        label="Contraste"
                    )
                    saturation = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=5,
                        step=1,
                        label="Saturation"
                    )

            # Paramètres de génération
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### ⚙️ Paramètres de Génération")
                with gr.Row():
                    quality = gr.Slider(
                        minimum=30,
                        maximum=50,
                        value=35,
                        label="Qualité"
                    )
                    creativity = gr.Slider(
                        minimum=5,
                        maximum=15,
                        value=7.5,
                        label="Créativité"
                    )

            with gr.Row():
                generate_btn = gr.Button("✨ Générer", variant="primary")
                clear_btn = gr.Button("🗑️ Effacer", variant="secondary")

            # Zone de résultat
            image_output = gr.Image(label="Aperçu")
            status = gr.Textbox(label="Statut", interactive=False)

        # Fonction de génération
        def generate(*args):
            logger.info("Démarrage d'une nouvelle génération")
            params = {
                "format_size": args[0],
                "orientation": args[1],
                "style": args[2],
                "layout": args[3],
                "ambiance": args[4],
                "palette": args[5],
                "subject": args[6],
                "title": args[7],
                "detail_level": args[8],
                "contrast": args[9],
                "saturation": args[10],
                "quality": args[11],
                "creativity": args[12]
            }
            result = generator.generate(params)
            logger.info(f"Génération terminée avec statut: {result[1]}")
            return result

        # Connexion des événements
        generate_btn.click(
            generate,
            inputs=[
                format_size,
                orientation,
                style,
                layout,
                ambiance,
                palette,
                subject,
                title,
                detail_level,
                contrast,
                saturation,
                quality,
                creativity
            ],
            outputs=[image_output, status]
        )

        clear_btn.click(
            lambda: (None, "🗑️ Image effacée"),
            outputs=[image_output, status]
        )

    logger.info("Interface créée avec succès")
    return app

if __name__ == "__main__":
    app = create_interface()
    app.launch()