import streamlit as st
import google.generativeai as genai
import os
from PIL import Image
import io

# Configuratiovbn de la page Streamlit
st.set_page_config(page_title="Mariam eco", layout="wide")
st.title("Mariam eco")
st.markdown("yup")

# Configuration de l'API Gemini
api_key = os.environ["GEMINI_API_KEY"] 

genai.configure(api_key=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éfinition du modèle
model = genai.GenerativeModel("gemini-2.5-flash-preview-04-17",system_instruction="Tu es Mariam une IA entraîner par Aenir pour répondre aux questions. ", safety_settings=safety_settings)

# Chargement de tous les fichiers PDF du dossier "data"
@st.cache_resource
def load_pdf_files():
    DATA_FOLDER = "data"
    if not os.path.exists(DATA_FOLDER):
        os.makedirs(DATA_FOLDER)
        st.warning(f"Le dossier {DATA_FOLDER} a été créé. Veuillez y placer vos fichiers PDF.")
        return []
    
    pdf_files = [
        os.path.join(DATA_FOLDER, file)
        for file in os.listdir(DATA_FOLDER)
        if file.lower().endswith(".pdf")
    ]
    
    if not pdf_files:
        st.info(f"Aucun fichier PDF trouvé dans le dossier {DATA_FOLDER}.")
        return []
    
    with st.spinner("Chargement des fichiers PDF..."):
        uploaded_files = [genai.upload_file(pdf_file) for pdf_file in pdf_files]
        st.success(f"{len(uploaded_files)} fichiers PDF chargés avec succès.")
    
    return uploaded_files

uploaded_files = load_pdf_files()

def process_input(prompt, image_file):
    """Traite le texte et l'image optionnelle et génère une réponse."""
    try:
        # Préparer le contenu en ajoutant tous les fichiers PDF
        print(prompt)
        content = uploaded_files.copy()

        # Si une image est fournie, l'ajouter au contenu
        if image_file is not None:
            # Streamlit fournit un fichier en mémoire, nous devons le sauvegarder temporairement
            temp_image_path = "temp_image.jpg"
            img = Image.open(image_file)
            img.save(temp_image_path)
            img_file = genai.upload_file(temp_image_path)
            content.append(img_file)
            # Supprimer le fichier temporaire
            os.remove(temp_image_path)

        # Ajout du prompt avec l'instruction de répondre en français
        content.append("\n\n")
        content.append(prompt + " répond en français.")

        # Génération du contenu
        with st.spinner("Génération de la réponse..."):
            result = model.generate_content(content)
        return result.text
    except Exception as e:
        return f"Une erreur s'est produite : "

# Interface principale
tab1, tab2 = st.tabs(["Interface principale", "Exemples"])

with tab1:
    st.subheader("Posez une question")
    
    # Zone de texte pour la question
    prompt = st.text_area("Question", placeholder="Posez une question sur le roman...", height=100)
    
    # Upload d'image (facultatif)
    image_file = st.file_uploader("Image (facultative)", type=["jpg", "jpeg", "png"])
    
    # Affichage de l'image téléchargée
    if image_file is not None:
        st.image(image_file, caption="Image téléchargée", use_column_width=True)
    
    # Bouton pour soumettre la question
    if st.button("Obtenir une réponse"):
        if prompt:
            response = process_input(prompt, image_file)
            st.markdown("### Réponse")
            st.markdown(response)
        else:
            st.warning("Veuillez entrer une question.")

with tab2:
    st.subheader("Exemples de questions")
    
    example1 = st.button("C'est quoi un mass media?")
    example2 = st.button("C'est quoi l'entrepreneuriat?")
    
    if example1:
        response = process_input("C'est quoi un mass media", None)
        st.markdown("### Réponse")
        st.markdown(response)
    
    if example2:
        response = process_input("C'est quoi l'entrepreneuriat", None)
        st.markdown("### Réponse")
        st.markdown(response)

# Affichage des informations sur les fichiers chargés
with st.sidebar:
    st.subheader("Informations")
    st.write(f"Nombre de fichiers PDF chargés: {len(uploaded_files)}")
    if uploaded_files:
        st.write("Fichiers PDF dans le dossier 'data':")
        for file in os.listdir("data"):
            if file.lower().endswith(".pdf"):
                st.write(f"- {file}")