import gradio as gr
import random
from PIL import Image, ImageDraw, ImageFont
import requests
from io import BytesIO
import os
import time

# Configuration
HF_TOKEN = "hf_..."  # Remplacez par votre token Hugging Face
API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"

STYLES = {
    "signage": "Signalétique Professionnelle",
    "poster": "Affiche Artistique",
    "silhouette": "Silhouette Décorative"
}

FORMATS = ["A4", "A3", "A2", "A1"]
MATERIALS = ["Lisse", "Ondulé", "Kraft"]

PROMPT_SUGGESTIONS = {
    "signage": [
        "Panneau directionnel moderne en carton recyclé, style minimaliste",
        "Présentoir écologique pour boutique, motifs naturels",
        "Signalétique murale en carton kraft, design épuré"
    ],
    "poster": [
        "Affiche artistique sur carton texturé, thème environnemental",
        "Poster graphique moderne, motifs géométriques en relief",
        "Design abstrait sur carton ondulé, couleurs naturelles"
    ],
    "silhouette": [
        "Découpe décorative florale en carton, style organique",
        "Silhouette architecturale moderne en carton recyclé",
        "Motif géométrique découpé, effet d'ombre et lumière"
    ]
}

STYLE_PROMPTS = {
    "signage": "professional signage design, minimalist, clean, photorealistic",
    "poster": "artistic poster design, creative, expressive, highly detailed",
    "silhouette": "decorative silhouette, elegant cutout design, clear shape"
}

def create_error_image(message):
    """Crée une image d'erreur avec un message."""
    img = Image.new('RGB', (512, 512), color='lightgray')
    d = ImageDraw.Draw(img)
    font = ImageFont.load_default()
    d.text((10,10), message, fill=(0,0,0), font=font)
    return img

def generate_image(prompt, style, format_size, material):
    """Génère une image via l'API Hugging Face avec gestion des erreurs et retries."""
    enhanced_prompt = f"a {material} cardboard {style} design showing {prompt}, {STYLE_PROMPTS[style]}, professional photography, 8k uhd, detailed"
    
    payload = {
        "inputs": enhanced_prompt,
        "negative_prompt": "low quality, blurry, bad anatomy, distorted, disfigured, pixelated",
        "num_inference_steps": 40,
        "guidance_scale": 7.5,
        "width": 768,
        "height": 768
    }

    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    max_retries = 3
    for attempt in range(max_retries):
        try:
            response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
            response.raise_for_status()
            return Image.open(BytesIO(response.content))
        except requests.exceptions.RequestException as e:
            if attempt == max_retries - 1:
                print(f"Erreur API après {max_retries} tentatives: {str(e)}")
                return create_error_image(f"Erreur API: {str(e)}")
            time.sleep(2 ** attempt)  # Attente exponentielle entre les tentatives

def update_suggestions(style):
    """Met à jour les suggestions de prompts selon le style choisi."""
    return gr.Dropdown(choices=PROMPT_SUGGESTIONS.get(style, []))

def apply_watermark(image):
    """Applique un filigrane à l'image générée."""
    draw = ImageDraw.Draw(image)
    font = ImageFont.load_default()
    text = "Equity Creation Studio"
    textwidth, textheight = draw.textsize(text, font)
    width, height = image.size
    x = width - textwidth - 10
    y = height - textheight - 10
    draw.text((x, y), text, font=font, fill=(255, 255, 255, 128))
    return image

def generate_and_process(prompt, style, format_size, material):
    """Génère l'image et applique le post-traitement."""
    image = generate_image(prompt, style, format_size, material)
    return apply_watermark(image)

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎨 Equity Creation Studio
        ### Studio de Design Éco-responsable
        Transformez vos idées en créations durables avec notre générateur d'images IA
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            style = gr.Radio(
                choices=list(STYLES.keys()),
                value="signage",
                label="Style de Création",
                info="Choisissez le type de design souhaité"
            )
            
            format_size = gr.Radio(
                choices=FORMATS,
                value="A4",
                label="Format",
                info="Sélectionnez la taille de votre création"
            )
            
            material = gr.Radio(
                choices=MATERIALS,
                value="Lisse",
                label="Type de Carton",
                info="Choisissez votre matériau"
            )
            
            prompt = gr.Textbox(
                lines=3,
                label="Description",
                placeholder="Décrivez votre vision..."
            )
            
            suggestions = gr.Dropdown(
                choices=PROMPT_SUGGESTIONS["signage"],
                label="Suggestions de prompts",
                info="Sélectionnez une suggestion ou écrivez votre propre description"
            )
            
            generate_btn = gr.Button("🪄 Générer", variant="primary")

        with gr.Column(scale=2):
            output = gr.Image(label="Image générée")

    style.change(fn=update_suggestions, inputs=[style], outputs=[suggestions])
    suggestions.change(fn=lambda x: x, inputs=[suggestions], outputs=[prompt])
    generate_btn.click(
        fn=generate_and_process,
        inputs=[prompt, style, format_size, material],
        outputs=[output]
    )

    gr.Markdown(
        """
        ### 🌿 Design Éco-responsable
        Equity Creation Studio transforme vos déchets carton en œuvres d'art
        """
    )

if __name__ == "__main__":
    demo.launch()