import gradio as gr
import google.generativeai as genai
import os
from PIL import Image

# Configurez l'API Gemini
genai.configure(api_key=os.environ["GEMINI_API_KEY"])

safety_settings = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
]

# Définissez le modèle
model = genai.GenerativeModel("gemini-2.0-flash-exp", safety_settings=safety_settings)

# Chemin vers le fichier PDF (assurez-vous qu'il est correct)
PDF_FILE_PATH = "okoumba-nkoghe-elo_compress.pdf"

# Uploadez le fichier PDF
sample_file = genai.upload_file(PDF_FILE_PATH)

def process_input(prompt, image_file=None):
    """
    Traite le fichier PDF et/ou l'image, et génère une réponse basée sur l'invite.
    
    :param prompt: La question ou l'invite de l'utilisateur
    :param pdf_file: Le fichier PDF uploadé (optionnel)
    :param image_file: L'image uploadée (optionnelle)
    :return: La réponse générée
    """
    try:
        # Préparez les entrées pour le modèle
        contents = []
        
        # Ajoutez le PDF s'il est présent
        if sample_file:
            contents.append(sample_file)
            contents.append("\n\n")
        
        # Ajoutez l'image si elle est uploadée
        if image_file:
            # Convertissez l'image de Gradio en objet PIL
            pil_image = Image.fromarray(image_file)
            # Convertissez l'image PIL en format compatible avec Gemini
            gemini_image = genai.upload_file(pil_image)
            contents.append(gemini_image)
            contents.append("\n\n")
        
        # Ajoutez le texte de l'invite
        contents.append(prompt + " répond en français.")
        
        # Générez le contenu
        result = model.generate_content(contents)
        
        return result.text
    except Exception as e:
        return f"Une erreur s'est produite : {str(e)}"

# Créez l'interface Gradio avec upload d'images optionnel
iface = gr.Interface(
    fn=process_input,
    inputs=[
        gr.Textbox(label="Question", placeholder="Posez une question sur le roman...", lines=4),
        #gr.File(label="PDF", type="file", visible=False),  # PDF est déjà uploadé, donc caché
        gr.Image(label="Image (optionnelle)", type="numpy", optional=True)
    ],
    outputs=gr.Textbox(label="Résultat"),
    title="Mariam elo la fille du soleil",
    description="Triché !! on va encore triché ohh !!!!"
)

# Lancez l'application
iface.launch()