import streamlit as st
import google.generativeai as genai
import os
from dotenv import load_dotenv
import http.client
import json
import shutil

# 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"},
]

# Instructions système pour Mariam (inchangées pour brièveté, mais peuvent être affinées)
ss = """
# Prompt System pour Mariam, IA conçu par youssouf 
[...]
"""

# Initialisation du modèle
model = genai.GenerativeModel('gemini-2.0-flash-exp', tools='code_execution',
                              safety_settings=safety_settings,
                              system_instruction=ss)

# Fonction de recherche web
def perform_web_search(query):
    conn = http.client.HTTPSConnection("google.serper.dev")
    payload = json.dumps({"q": query})
    headers = {
        'X-API-KEY': '9b90a274d9e704ff5b21c0367f9ae1161779b573',
        'Content-Type': 'application/json'
    }
    try:
        conn.request("POST", "/search", payload, headers)
        res = conn.getresponse()
        data = json.loads(res.read().decode("utf-8"))
        return data
    except Exception as e:
        st.error(f"Erreur lors de la recherche web : {e}")
        return None
    finally:
        conn.close()

# Formater les résultats de recherche
def format_search_results(data):
    if not data:
        return "Aucun résultat trouvé"
    result = ""
    if 'knowledgeGraph' in data:
        kg = data['knowledgeGraph']
        result += f"### {kg.get('title', '')}\n*{kg.get('type', '')}*\n\n{kg.get('description', '')}\n\n"
    if 'organic' in data:
        result += "### Résultats principaux:\n"
        for item in data['organic'][:3]:
            result += f"- **{item['title']}**\n  {item['snippet']}\n  [Lien]({item['link']})\n\n"
    if 'peopleAlsoAsk' in data:
        result += "### Questions fréquentes:\n"
        for item in data['peopleAlsoAsk'][:2]:
            result += f"- **{item['question']}**\n  {item['snippet']}\n\n"
    return result

# Convertir les rôles pour Streamlit
def role_to_streamlit(role):
    return "assistant" if role == "model" else role

# Initialiser l'état de la session
if "chat" not in st.session_state:
    st.session_state.chat = model.start_chat(history=[])
if "web_search" not in st.session_state:
    st.session_state.web_search = False
if "messages" not in st.session_state:
    st.session_state.messages = []

# Fonction pour effacer l’historique
def clear_chat_history():
    st.session_state.chat = model.start_chat(history=[])
    st.session_state.messages = []

# Créer un dossier temporaire
os.makedirs("temp", exist_ok=True)

# Titre de l’application
st.title("Mariam AI")

# Section des paramètres dans la barre latérale
with st.sidebar:
    st.header("Paramètres")
    st.session_state.web_search = st.toggle("Activer la recherche web", value=st.session_state.web_search, 
                                            help="Permet d’enrichir les réponses avec des informations du web.")
    if st.button("Effacer l’historique", on_click=clear_chat_history):
        st.success("Historique effacé.")

# Section de téléchargement de fichiers
with st.sidebar:
    st.header("Téléchargement de fichiers")
    st.info("Types acceptés : jpg, jpeg, png, pdf, txt")
    uploaded_file = st.file_uploader("Choisir un fichier", type=['jpg', 'jpeg', 'png', 'pdf', 'txt'])

# Section de conversation
st.header("Conversation")
chat_container = st.container()

# Afficher les messages
with chat_container:
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

# Champ de saisie multiligne
user_input = st.text_area("Votre message", height=100)

# Bouton pour envoyer
if st.button("Envoyer"):
    if user_input:
        # Ajouter le message de l’utilisateur
        st.session_state.messages.append({"role": "user", "content": user_input})
        
        # Traiter le fichier téléchargé
        uploaded_gemini_file = None
        if uploaded_file:
            file_ext = uploaded_file.name.split('.')[-1].lower()
            accepted_types = ['jpg', 'jpeg', 'png', 'pdf', 'txt']
            if file_ext in accepted_types:
                with open(os.path.join("temp", uploaded_file.name), "wb") as f:
                    f.write(uploaded_file.getbuffer())
                try:
                    uploaded_gemini_file = genai.upload_file(os.path.join("temp", uploaded_file.name))
                except Exception as e:
                    st.error(f"Erreur lors de l’upload du fichier : {e}")
            else:
                st.error("Type de fichier non accepté. Utilisez jpg, jpeg, png, pdf ou txt.")
        
        # Recherche web si activée
        web_results = None
        if st.session_state.web_search:
            with st.spinner("Recherche web en cours..."):
                web_results = perform_web_search(user_input)
                if web_results:
                    formatted_results = format_search_results(web_results)
                    user_input = f"{user_input}\n\nVoici les résultats de la recherche web. Analyse-les et donne-moi une réponse complète :\n\n{formatted_results}"
        
        # Envoyer à Gemini
        try:
            if uploaded_gemini_file:
                response = st.session_state.chat.send_message([uploaded_gemini_file, "\n\n", user_input])
            else:
                response = st.session_state.chat.send_message(user_input)
            st.session_state.messages.append({"role": "assistant", "content": response.text})
            with chat_container:
                with st.chat_message("assistant"):
                    st.markdown(response.text)
        except Exception as e:
            st.error(f"Erreur lors de l’envoi : {e}")
        
        # Nettoyer les fichiers temporaires
        if uploaded_file:
            shutil.rmtree("temp", ignore_errors=True)
            os.makedirs("temp", exist_ok=True)