import streamlit as st
import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

# Configuration de la page
st.set_page_config(
    page_title="Locgi - Chatbot",
    page_icon="💬",
    layout="centered"
)

# Titre et description
st.title("💬 Locgi - Votre Assistant")
st.markdown("""
Bienvenue ! Je suis Locgi, votre assistant conversationnel. 
Posez-moi vos questions et discutons ensemble !
""")

# Initialisation de l'historique de conversation
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "assistant", "content": "Bonjour ! Je suis Locgi. Comment puis-je vous aider aujourd'hui ?"}
    ]

# Initialisation du modèle de conversation
@st.cache_resource
def load_model():
    try:
        # Vérifier si CUDA est disponible
        device = "cuda" if torch.cuda.is_available() else "cpu"
        
        # Charger le modèle et le tokenizer
        model_name = "openai/gpt-oss-20b"
        
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            low_cpu_mem_usage=True
        )
        
        # Déplacer le modèle sur le bon device
        model = model.to(device)
        
        # Créer le pipeline
        generator = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            device=0 if device == "cuda" else -1
        )
        
        return generator, device
    except Exception as e:
        st.error(f"Erreur lors du chargement du modèle : {e}")
        return None, None

# Chargement du modèle
chatbot, device = load_model()

# Affichage de l'information sur le device
if device:
    st.sidebar.info(f"Modèle chargé sur : {'GPU' if device == 'cuda' else 'CPU'}")
else:
    st.sidebar.warning("Modèle non chargé")

# Affichage de l'historique des messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Zone de saisie utilisateur
if prompt := st.chat_input("Écrivez votre message ici..."):
    # Ajout du message utilisateur à l'historique
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Affichage du message utilisateur
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Génération de la réponse du bot
    if chatbot:
        with st.chat_message("assistant"):
            with st.spinner("Locgi réfléchit..."):
                try:
                    # Construire le contexte de conversation
                    conversation_history = ""
                    for msg in st.session_state.messages[-5:]:  # Garder les 5 derniers messages pour le contexte
                        if msg["role"] == "user":
                            conversation_history += f"Utilisateur: {msg['content']}\n"
                        else:
                            conversation_history += f"Locgi: {msg['content']}\n"
                    
                    conversation_history += f"Locgi: "
                    
                    # Génération de la réponse
                    response = chatbot(
                        conversation_history,
                        max_new_tokens=150,
                        do_sample=True,
                        temperature=0.7,
                        top_p=0.9,
                        pad_token_id=chatbot.tokenizer.eos_token_id
                    )[0]['generated_text']
                    
                    # Extraire la réponse générée
                    if "Locgi: " in response:
                        response = response.split("Locgi: ")[-1].strip()
                    
                    # Limiter la réponse à la première phrase complète si nécessaire
                    if len(response) > 0:
                        st.markdown(response)
                    else:
                        response = "Désolé, je n'ai pas compris. Pouvez-vous reformuler ?"
                        st.markdown(response)
                    
                    # Ajout de la réponse à l'historique
                    st.session_state.messages.append(
                        {"role": "assistant", "content": response}
                    )
                except Exception as e:
                    error_msg = f"Désolé, une erreur est survenue : {str(e)}"
                    st.error(error_msg)
                    st.session_state.messages.append(
                        {"role": "assistant", "content": "Désolé, j'ai rencontré un problème technique."}
                    )
    else:
        st.error("Le modèle n'a pas pu être chargé. Veuillez réessayer ou utiliser un autre modèle.")

# Bouton pour effacer la conversation
if st.sidebar.button("🗑️ Effacer la conversation"):
    st.session_state.messages = [
        {"role": "assistant", "content": "Bonjour ! Je suis Locgi. Comment puis-je vous aider aujourd'hui ?"}
    ]
    st.rerun()

# Paramètres de génération dans la barre latérale
st.sidebar.markdown("---")
st.sidebar.markdown("### Paramètres de génération")
temperature = st.sidebar.slider("Température", 0.1, 1.5, 0.7, 0.1)
max_tokens = st.sidebar.slider("Tokens maximum", 50, 500, 150, 10)

# Pied de page
st.sidebar.markdown("---")
st.sidebar.markdown("### À propos de Locgi")
st.sidebar.markdown("Locgi est un assistant conversationnel basé sur le modèle GPT d'OpenAI.")