import os
import gradio as gr
from huggingface_hub import hf_hub_download, login
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoPipeline
from pptx import Presentation
from pptx.util import Inches, Pt
from pptx.enum.text import PP_ALIGN
import torch
from llama_cpp import Llama
import time
from PIL import Image
import io
import requests

# Configuration des modèles disponibles
TEXT_MODELS = {
    "Mistral Small 24B (GGUF)": "MisterAI/Bartowski_MistralAI_Mistral-Small-24B-Base-2501-GGUF",
    "Mixtral 8x7B": "mistralai/Mixtral-8x7B-v0.1",
    "Lucie 7B": "OpenLLM-France/Lucie-7B"
}

IMAGE_MODELS = {
    "FLUX.1": "black-forest-labs/FLUX.1-schnell",
    "ArtifyAI": "ImageInception/ArtifyAI-v1.1"
}

# Préprompt amélioré pour une meilleure structuration
PREPROMPT = """Vous êtes un assistant IA expert en création de présentations PowerPoint professionnelles. 
Générez une présentation structurée et détaillée en suivant ce format EXACT:

TITRE: [Titre principal de la présentation]

DIAPO 1:
Titre: [Titre de la diapo]
Points:
- Point 1
- Point 2
- Point 3
Image: [Description détaillée de l'image souhaitée pour cette diapo]

DIAPO 2:
Titre: [Titre de la diapo]
Points:
- Point 1
- Point 2
- Point 3
Image: [Description détaillée de l'image souhaitée pour cette diapo]

[Continuez avec ce format pour chaque diapositive]

Analysez le texte suivant et créez une présentation professionnelle avec des descriptions d'images pertinentes :"""

class PresentationGenerator:
    def __init__(self):
        self.token = os.getenv('Authentification_HF')
        if not self.token:
            raise ValueError("Token d'authentification HuggingFace non trouvé")
        login(self.token)
        self.text_model = None
        self.text_tokenizer = None
        self.image_pipeline = None
        
    def load_text_model(self, model_name):
        """Charge le modèle de génération de texte"""
        model_id = TEXT_MODELS[model_name]
        if model_id.endswith('.gguf'):
            # Configuration pour les modèles GGUF
            model_path = hf_hub_download(
                repo_id=model_id.split('/')[0] + '/' + model_id.split('/')[1],
                filename=model_id.split('/')[-1],
                token=self.token
            )
            self.text_model = Llama(
                model_path=model_path,
                n_ctx=4096,
                n_batch=512,
                verbose=False
            )
        else:
            # Configuration pour les modèles Transformers standards
            self.text_tokenizer = AutoTokenizer.from_pretrained(model_id, token=self.token)
            self.text_model = AutoModelForCausalLM.from_pretrained(
                model_id,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                token=self.token
            )
        
    def load_image_model(self, model_name):
        """Charge le modèle de génération d'images"""
        model_id = IMAGE_MODELS[model_name]
        self.image_pipeline = AutoPipeline.from_pretrained(
            model_id,
            token=self.token
        )

    def generate_text(self, prompt, temperature=0.7, max_tokens=4096):
        """Génère le texte de la présentation"""
        if isinstance(self.text_model, Llama):
            response = self.text_model(
                prompt,
                max_tokens=max_tokens,
                temperature=temperature,
                echo=False
            )
            return response['choices'][0]['text']
        else:
            inputs = self.text_tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}],
                return_tensors="pt",
                return_dict=True
            )
            outputs = self.text_model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                temperature=temperature
            )
            return self.text_tokenizer.decode(outputs[0], skip_special_tokens=True)

    def generate_image(self, prompt, negative_prompt="", num_inference_steps=30):
        """Génère une image pour la diapositive"""
        try:
            image = self.image_pipeline(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=num_inference_steps
            ).images[0]
            return image
        except Exception as e:
            print(f"Erreur lors de la génération de l'image: {str(e)}")
            return None

    def parse_presentation_content(self, content):
        """Parse le contenu généré en sections pour les diapositives"""
        slides = []
        current_slide = None
        
        for line in content.split('\n'):
            line = line.strip()
            if line.startswith('TITRE:'):
                slides.append({'type': 'title', 'title': line[6:].strip()})
            elif line.startswith('DIAPO'):
                if current_slide:
                    slides.append(current_slide)
                current_slide = {'type': 'content', 'title': '', 'points': [], 'image_prompt': ''}
            elif line.startswith('Titre:') and current_slide:
                current_slide['title'] = line[6:].strip()
            elif line.startswith('- ') and current_slide:
                current_slide['points'].append(line[2:].strip())
            elif line.startswith('Image:') and current_slide:
                current_slide['image_prompt'] = line[6:].strip()
        
        if current_slide:
            slides.append(current_slide)
        
        return slides

    def create_presentation(self, slides):
        """Crée la présentation PowerPoint avec texte et images"""
        prs = Presentation()
        
        # Première diapo (titre)
        title_slide = prs.slides.add_slide(prs.slide_layouts[0])
        title_slide.shapes.title.text = slides[0]['title']
        
        # Autres diapos
        for slide in slides[1:]:
            content_slide = prs.slides.add_slide(prs.slide_layouts[1])
            content_slide.shapes.title.text = slide['title']
            
            # Ajout du texte
            if slide['points']:
                body = content_slide.shapes.placeholders[1].text_frame
                body.clear()
                for point in slide['points']:
                    p = body.add_paragraph()
                    p.text = point
                    p.level = 0
            
            # Ajout de l'image si disponible
            if slide.get('image_prompt'):
                image = self.generate_image(slide['image_prompt'])
                if image:
                    # Sauvegarde temporaire de l'image
                    img_path = f"temp_slide_{slides.index(slide)}.png"
                    image.save(img_path)
                    
                    # Ajout de l'image à la diapositive
                    left = Inches(1)
                    top = Inches(2.5)
                    content_slide.shapes.add_picture(img_path, left, top, height=Inches(4))
                    
                    # Suppression du fichier temporaire
                    os.remove(img_path)
        
        return prs

def generate_presentation_with_progress(text, text_model_name, image_model_name, temperature, max_tokens, negative_prompt):
    """Fonction principale de génération avec suivi de progression"""
    try:
        start_time = time.time()
        generator = PresentationGenerator()
        
        # Chargement des modèles
        yield "Chargement des modèles...", None, None
        generator.load_text_model(text_model_name)
        generator.load_image_model(image_model_name)
        
        # Génération du contenu
        yield "Génération du contenu de la présentation...", None, None
        full_prompt = PREPROMPT + "\n\n" + text
        generated_content = generator.generate_text(full_prompt, temperature, max_tokens)
        
        # Création de la présentation
        yield "Création de la présentation PowerPoint...", generated_content, None
        slides = generator.parse_presentation_content(generated_content)
        prs = generator.create_presentation(slides)
        
        # Sauvegarde
        output_path = "presentation.pptx"
        prs.save(output_path)
        
        execution_time = time.time() - start_time
        status = f"Présentation générée avec succès en {execution_time:.2f} secondes!"
        
        return status, generated_content, output_path
        
    except Exception as e:
        return f"Erreur: {str(e)}", None, None

# Interface Gradio améliorée
css = """
/* Thème sombre personnalisé */
.gradio-container {
    background-color: #000000 !important;
}

.gr-form, .gr-box, .gr-panel {
    border-radius: 8px !important;
    background-color: #1a1a1a !important;
    border: 1px solid #333333 !important;
}

.gr-input, .gr-textarea, .gr-dropdown {
    background-color: #2d2d2d !important;
    color: #ffffff !important;
    border: 1px solid #404040 !important;
}

.gr-button {
    background-color: #2d2d2d !important;
    color: #ffffff !important;
    border: 1px solid #404040 !important;
    transition: all 0.3s ease !important;
}

.gr-button:hover {
    background-color: #404040 !important;
    transform: translateY(-2px) !important;
}

/* Textes et labels */
h1, h2, h3, p, label, .gr-text {
    color: #ffffff !important;
}

/* Scrollbar */
::-webkit-scrollbar {
    width: 8px;
    height: 8px;
}

::-webkit-scrollbar-track {
    background: #1a1a1a;
}

::-webkit-scrollbar-thumb {
    background: #404040;
    border-radius: 4px;
}

::-webkit-scrollbar-thumb:hover {
    background: #4a4a4a;
}
"""

with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
    gr.Markdown(
        """
        # 🎯 Générateur de Présentations PowerPoint IA
        
        Créez des présentations professionnelles automatiquement avec l'aide de l'IA.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            text_model_choice = gr.Dropdown(
                choices=list(TEXT_MODELS.keys()),
                value=list(TEXT_MODELS.keys())[0],
                label="Modèle de génération de texte"
            )
            image_model_choice = gr.Dropdown(
                choices=list(IMAGE_MODELS.keys()),
                value=list(IMAGE_MODELS.keys())[0],
                label="Modèle de génération d'images"
            )
            temperature = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.7,
                step=0.1,
                label="Température"
            )
            max_tokens = gr.Slider(
                minimum=1000,
                maximum=4096,
                value=2048,
                step=256,
                label="Tokens maximum"
            )
            negative_prompt = gr.Textbox(
                lines=2,
                label="Prompt négatif pour les images",
                placeholder="Ce que vous ne voulez pas voir dans les images..."
            )
    
    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                lines=10,
                label="Votre texte",
                placeholder="Décrivez le contenu que vous souhaitez pour votre présentation..."
            )
            file_upload = gr.File(
                label="Documents de référence (PDF, Images)",
                file_types=["pdf", "png", "jpg", "jpeg"],
                multiple=True
            )
    
    with gr.Row():
        generate_btn = gr.Button("🚀 Générer la présentation", variant="primary")
    
    with gr.Row():
        with gr.Column():
            status_output = gr.Textbox(
                label="Statut",
                lines=2
            )
            generated_content = gr.Textbox(
                label="Contenu généré",
                lines=10,
                show_copy_button=True
            )
            output_file = gr.File(
                label="Présentation PowerPoint"
            )
    
    generate_btn.click(
        fn=generate_presentation_with_progress,
        inputs=[
            input_text,
            text_model_choice,
            image_model_choice,
            temperature,
            max_tokens,
            negative_prompt
        ],
        outputs=[
            status_output,
            generated_content,
            output_file
        ]
    )

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