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

# Charger les variables d'environnement
load_dotenv()

# Configurer la clé API
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Paramètres de sécurité
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"},
]

# Initialiser le modèle
model = genai.GenerativeModel('gemini-1.5-flash', 
                               safety_settings=safety_settings,
                               system_instruction="Tu es un assistant intelligent. ton but est d'assister au mieux que tu peux. tu as été créé par Aenir et tu t'appelles Mariam")

# Fonction pour convertir le rôle
def role_to_streamlit(role):
    return "assistant" if role == "model" else role

# Configuration de la page Streamlit
st.set_page_config(page_title="Mariam - Assistant IA", page_icon="🤖")
st.title("Mariam AI - Chat Intelligent")

# Initialiser l'historique de chat si pas déjà présent
if "chat" not in st.session_state:
    st.session_state.chat = model.start_chat(history=[])

# Gestion de l'upload d'images
uploaded_file = st.file_uploader("Télécharger une image (optionnel)", 
                                  type=["jpg", "jpeg", "png", "gif", "bmp"])

# Afficher l'historique des messages
for message in st.session_state.chat.history:
    with st.chat_message(role_to_streamlit(message.role)):
        st.markdown(message.parts[0].text)
        # Afficher l'image si présente dans l'historique
        if len(message.parts) > 1 and hasattr(message.parts[1], 'image'):
            st.image(message.parts[1].image)

# Zone de saisie du chat
if prompt := st.chat_input("Que puis-je faire pour vous ?"):
    # Préparer le message
    if uploaded_file is not None:
        # Convertir le fichier téléchargé en image
        image = Image.open(uploaded_file)
        
        # Envoyer le message avec l'image
        response = st.session_state.chat.send_message([
            prompt, 
            image
        ])
        
        # Afficher l'image téléchargée
        st.chat_message("user").markdown(prompt)
        st.chat_message("user").image(image)
    else:
        # Envoyer le message texte seul
        response = st.session_state.chat.send_message(prompt)
        st.chat_message("user").markdown(prompt)

    # Afficher la réponse de l'assistant
    with st.chat_message("assistant"):
        st.markdown(response.text)