import gradio as gr
import os
from PIL import Image
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 Classiques
    "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"
    },
    "Surréaliste": {
        "prompt_prefix": "surrealist art, dreamlike imagery, symbolic elements, dali inspired, metaphysical art",
        "negative_prompt": "realistic, conventional, ordinary, literal"
    },

    # Styles Modernes
    "Art Moderne": {
        "prompt_prefix": "modern art style poster, professional design, contemporary aesthetic",
        "negative_prompt": "traditional, cluttered, busy design, vintage"
    },
    "Pop Art": {
        "prompt_prefix": "pop art style poster, bold colors, repeated patterns, screen print effect, warhol inspired",
        "negative_prompt": "subtle, realistic, traditional, painterly"
    },
    "Minimaliste": {
        "prompt_prefix": "minimalist design poster, clean composition, elegant simplicity",
        "negative_prompt": "complex, detailed, ornate, busy, cluttered"
    },
    "Cubiste": {
        "prompt_prefix": "cubist art style, geometric fragmentation, multiple perspectives, picasso inspired",
        "negative_prompt": "realistic, single perspective, traditional, photographic"
    },
    "Futuriste": {
        "prompt_prefix": "futuristic art style, dynamic movement, technological elements, speed and motion",
        "negative_prompt": "static, traditional, classical, historical"
    },

    # Styles Spéciaux
    "Neo Vintage": {
        "prompt_prefix": "vintage style advertising poster, retro design, classic aesthetic",
        "negative_prompt": "modern, digital, contemporary style"
    },
    "Cyberpunk": {
        "prompt_prefix": "cyberpunk style poster, neon lights, futuristic design, high-tech aesthetic",
        "negative_prompt": "vintage, natural, rustic, traditional"
    },
    "Japonais": {
        "prompt_prefix": "japanese art style poster, ukiyo-e inspired design, traditional japanese aesthetic",
        "negative_prompt": "western, modern, photographic"
    },
    "Art Déco": {
        "prompt_prefix": "art deco style poster, geometric patterns, luxury design, 1920s aesthetic",
        "negative_prompt": "modern, minimalist, casual, contemporary"
    },
    "Symboliste": {
        "prompt_prefix": "symbolic art, decorative patterns, gold elements, mystical atmosphere, klimt inspired",
        "negative_prompt": "realistic, simple, plain, literal"
    }
}

# Paramètres de composition enrichis
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:
        """Construction de prompt améliorée avec plus de détails"""
        style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
        prompt = f"{style_info['prompt_prefix']}, {params['subject']}"

        # Ajout des paramètres de composition
        if params.get("layout"):
            prompt += f", {COMPOSITION_PARAMS['Layouts'][params['layout']]}"
        if params.get("ambiance"):
            prompt += f", {COMPOSITION_PARAMS['Ambiances'][params['ambiance']]}"
        if params.get("palette"):
            prompt += f", {COMPOSITION_PARAMS['Palette'][params['palette']]}"

        # Ajout des ajustements fins
        if params.get("detail_level"):
            detail_strength = params["detail_level"]
            prompt += f", {'highly detailed, intricate' if detail_strength > 7 else 'moderately detailed'}"

        if params.get("contrast"):
            contrast_strength = params["contrast"]
            prompt += f", {'high contrast, dramatic lighting' if contrast_strength > 7 else 'balanced contrast'}"

        if params.get("saturation"):
            saturation_strength = params["saturation"]
            prompt += f", {'vibrant saturated colors' if saturation_strength > 7 else 'subtle muted colors'}"

        if params.get("title"):
            prompt += f", with text saying '{params['title']}'"

        logger.debug(f"Prompt final: {prompt}")
        return prompt

    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)
            
            payload = {
                "inputs": prompt,
                "parameters": {
                    "negative_prompt": ART_STYLES[params["style"]]["negative_prompt"],
                    "num_inference_steps": min(int(35 * (params["quality"]/100)), 40),
                    "guidance_scale": min(7.5 * (params["creativity"]/10), 10.0),
                    "width": 768,
                    "height": 768 if params["orientation"] == "Portrait" else 512
                }
            }
            
            logger.debug(f"Payload: {json.dumps(payload, indent=2)}")

            response = requests.post(
                self.API_URL,
                headers=self.headers,
                json=payload,
                timeout=30
            )
            
            if response.status_code == 200:
                image = Image.open(io.BytesIO(response.content))
                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 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"):
            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"
                    )
            
            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"
                    )
            
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 📝 Contenu")
                subject = gr.Textbox(
                    label="Description",
                    placeholder="Décrivez votre vision..."
                )
                title = gr.Textbox(
                    label="Titre",
                    placeholder="Titre de l'affiche..."
                )
            
            with gr.Group(elem_classes="advanced-controls"):
                gr.Markdown("### 🎯 Ajustements Fins")
                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"
                    )
            
            with gr.Group(elem_classes="controls-group"):
                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")
            
            image_output = gr.Image(label="Aperçu")
            status = gr.Textbox(label="Statut", interactive=False)
        
        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
        
        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()
    logger.info("Démarrage de l'application")
    app.launch()
    # [Gardez tous vos imports existants]
# Ajoutez ces imports en haut du fichier
from pathlib import Path
import time
import json
from datetime import datetime

# [Gardez toutes vos configurations existantes (ART_STYLES, etc.)]

# Ajoutez cette classe juste avant la classe ImageGenerator
class GenerationManager:
    def __init__(self):
        self.save_dir = Path("generations")
        self.save_dir.mkdir(exist_ok=True)
        self.history_file = self.save_dir / "history.json"
        self.history = self.load_history()

    def load_history(self):
        if self.history_file.exists():
            try:
                with open(self.history_file, "r", encoding="utf-8") as f:
                    return json.load(f)
            except Exception as e:
                logger.error(f"Erreur chargement historique: {e}")
                return []
        return []

    def save_generation(self, image, params):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        image_name = f"generation_{timestamp}.png"
        image_path = self.save_dir / image_name

        # Sauvegarde de l'image
        image.save(image_path)

        # Création de l'entrée
        entry = {
            "id": timestamp,
            "image_path": str(image_path),
            "params": params,
            "created_at": timestamp
        }

        self.history.append(entry)
        self._save_history()
        return entry

    def _save_history(self):
        with open(self.history_file, "w", encoding="utf-8") as f:
            json.dump(self.history, f, indent=2, ensure_ascii=False)

    def get_recent_generations(self, limit=10):
        return sorted(
            self.history,
            key=lambda x: x["created_at"],
            reverse=True
        )[:limit]

# [Gardez votre classe ImageGenerator existante]

# Modifiez votre fonction create_interface comme ceci
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()
    generation_manager = GenerationManager()  # Ajout du gestionnaire
    
    with gr.Blocks(css=css) as app:
        # [Gardez toute votre interface existante jusqu'à la fin]
        
        # Ajoutez cet onglet à la fin, avant les boutons de génération
        with gr.Tab("📜 Historique"):
            with gr.Row():
                history_gallery = gr.Gallery(
                    label="Générations Récentes",
                    columns=3,
                    height=400,
                    show_label=True
                )
                
            with gr.Row():
                refresh_btn = gr.Button("🔄 Actualiser")
        
        # Modification de la fonction generate existante
        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)
            
            # Ajout de la sauvegarde automatique
            if result[0] is not None:
                generation_manager.save_generation(result[0], params)
            
            logger.info(f"Génération terminée avec statut: {result[1]}")
            return result

        def refresh_history():
            recent = generation_manager.get_recent_generations()
            return gr.Gallery.update(value=[entry["image_path"] for entry in recent])

        # Connexion des boutons
        generate_btn.click(
            generate,
            inputs=[
                format_size,
                orientation,
                style,
                layout,
                ambiance,
                palette,
                subject,
                title,
                detail_level,
                contrast,
                saturation,
                quality,
                creativity
            ],
            outputs=[image_output, status]
        )
        
        refresh_btn.click(
            refresh_history,
            outputs=[history_gallery]
        )
        
        clear_btn.click(
            lambda: (None, "🗑️ Image effacée"),
            outputs=[image_output, status]
        )
    
    logger.info("Interface créée avec succès")
    return app