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

load_dotenv()

# Configure the API key
genai.configure(api_key=os.getenv("GOOGLE_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"},
]

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

def role_to_streamlit(role):
    if role == "model":
        return "assistant"
    else:
        return role

# Add a Gemini Chat history object to Streamlit session state
if "chat" not in st.session_state:
    st.session_state.chat = model.start_chat(history=[])

# Display Form Title
st.title("Mariam AI!")

# Display chat messages from history above current input box
for message in st.session_state.chat.history:
    with st.chat_message(role_to_streamlit(message.role)):
        for part in message.parts:
            if part.text:  # Check for text content
                st.markdown(part.text)
            elif part.file_data:  # Check for file data
                try:
                    # Infer MIME type if not provided
                    mime_type = part.file_data.mime_type
                    if not mime_type:
                        mime_type = mimetypes.guess_type(part.file_data.file_name)[0]

                    if mime_type and mime_type.startswith("image/"):
                        image = Image.open(io.BytesIO(part.file_data.data))
                        st.image(image)
                    else:
                        st.write(f"File: {part.file_data.file_name} (MIME type: {mime_type})")
                except Exception as e:
                    st.error(f"Error displaying file: {e}")

# Accept user's next message and file uploads
if prompt := st.chat_input("Hey?"):
    uploaded_file = st.file_uploader("Choose a file", type=["jpg", "jpeg", "png", "pdf"])

    parts = [prompt]
    if uploaded_file:
        bytes_data = uploaded_file.getvalue()
        parts.append({
            "file_data": {
                "mime_type": uploaded_file.type,
                "file_name": uploaded_file.name,
                "data": bytes_data
            }
        })

        # Display the uploaded image
        if uploaded_file.type.startswith("image/"):
            image = Image.open(uploaded_file)
            with st.chat_message("user"):
                st.image(image, caption=f"Uploaded Image: {uploaded_file.name}")

    # Display user's message
    with st.chat_message("user"):
        st.markdown(prompt)

    # Send message to Gemini
    response = st.session_state.chat.send_message(parts)

    # Display Gemini's response
    with st.chat_message("assistant"):
        st.markdown(response.text)