import os
import gradio as gr
from huggingface_hub import hf_hub_download, login
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
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
from diffusers import FluxPipeline

# Configuration des modèles disponibles
TEXT_MODELS = {
    "Mistral Nemo 2407 (GGUF)": "MisterAI/Bartowski_MistralAI_Mistral-Nemo-Instruct-2407-IQ4_XS.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. Soyez très précis dans la description pour permettre
une génération d'image de qualité. Par exemple : "Une illustration professionnelle montrant un concept clé de cybersécurité
avec des éléments visuels modernes, un style épuré et des couleurs corporate (fond noir, couleurs bleu electrique, rouge, gris, blanc).
L'image doit être claire, minimaliste et adaptée à une présentation professionnelle."]

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]
        if model_id == "black-forest-labs/FLUX.1-schnell":
            self.image_pipeline = FluxPipeline.from_pretrained(
                model_id,
                revision="refs/pr/1",  # Utiliser une révision spécifique
                torch_dtype=torch.bfloat16
            )
            self.image_pipeline.enable_model_cpu_offload()  # Économise de la VRAM en déchargeant le modèle sur le CPU
            self.image_pipeline.tokenizer.add_prefix_space = False  # Désactive add_prefix_space
            print(f"Modèle d'image FLUX chargé : {model_id}")
        else:
            self.image_pipeline = pipeline(
                "text-to-image",
                model=model_id,
                token=self.token
            )
            print(f"Modèle d'image chargé : {model_id}")

    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,
                guidance_scale=0.0,
                max_sequence_length=256,
                generator=torch.Generator("cpu").manual_seed(0)
            ).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

# CSS personnalisé pour un thème sombre amélioré
css = """
/* Thème sombre personnalisé */
.gradio-container {
    background-color: #000000 !important;
    color: #ffffff !important;
}

.gr-form, .gr-box, .gr-panel {
    border-radius: 8px !important;
    background-color: #1a1a1a !important;
    border: 1px solid #333333 !important;
    color: #ffffff !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;
}

/* Assurez-vous que les éléments de formulaire et les boutons s'affichent correctement */
.gr-form input, .gr-form textarea, .gr-form select {
    background-color: #2d2d2d !important;
    color: #ffffff !important;
    border: 1px solid #404040 !important;
}

.gr-form label {
    color: #ffffff !important;
}

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

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

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..."
            )
            # Modification du composant File pour supprimer l'argument multiple
            file_upload = gr.File(
                label="Documents de référence (PDF, Images)",
                file_types=["pdf", "png", "jpg", "jpeg"]
            )

    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()