import streamlit as st
import os
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import login

# Read the Hugging Face API token from the environment variable
api_token = os.getenv("HF_API_TOKEN")
if not api_token:
    st.error("Hugging Face API token is missing. Please add it to the secrets in the Space settings.")
    st.stop()

# Login using your Hugging Face token
login(token=api_token)

model_id = "meta-llama/Meta-Llama-3-8B"

# Initialize the model pipeline with authentication
pipe = pipeline("text-generation", model=model_id, use_auth_token=api_token)

# Display the logo and title
st.image("logo.jpg", width=300)
st.title("Coach Virtual PRODI")

# Initialize a session state variable for history if it doesn't exist
if 'history' not in st.session_state:
    st.session_state['history'] = []

# Function to update the conversation history
def update_history(user_input, ai_response):
    st.session_state['history'].append(("User", user_input))
    st.session_state['history'].append(("AI", ai_response))

# Display the conversation history
for speaker, text in st.session_state['history']:
    if speaker == "User":
        st.text_input("Usuario", value=text, disabled=True)
    else:
        st.text_area("PRODI", value=text, height=75, disabled=True)

# Chat input for user prompt
user_input = st.text_input("¿Cómo te puedo ayudar hoy?")
if user_input:
    with st.spinner("Generando respuesta..."):
        # Get the AI's response
        ai_response = pipe(user_input, max_length=100, num_return_sequences=1)[0]['generated_text']
        
        # Update the conversation history
        update_history(user_input, ai_response)
        
        # Display the AI's response
        st.write(ai_response)