import os
import streamlit as st
from groq import Groq

# Streamlit page configuration
st.set_page_config(
    page_title="LLAMA 3.1 Chat",
    page_icon="🦙",
    layout="centered"
)

# Retrieve API key from Hugging Face Secrets
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

if not GROQ_API_KEY:
    st.error("⚠️ Error: GROQ_API_KEY is missing! Please add it to Hugging Face Secrets.")
    st.stop()

# Initialize Groq client with API key
client = Groq(api_key=GROQ_API_KEY)

# Initialize the chat history in Streamlit session state if not present already
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Streamlit page title
st.title("🦙 LLAMA 3.1 ChatBot")

# Display chat history
for message in st.session_state.chat_history:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Input field for user's message
user_prompt = st.chat_input("Ask LLAMA...")

if user_prompt:
    st.chat_message("user").markdown(user_prompt)
    st.session_state.chat_history.append({"role": "user", "content": user_prompt})

    # Send user's message to the LLM and get a response
    messages = [
        {"role": "system", "content": "You are a helpful assistant"},
        *st.session_state.chat_history
    ]

    try:
        response = client.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=messages
        )

        assistant_response = response.choices[0].message.content
        st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})

        # Display the LLM's response
        with st.chat_message("assistant"):
            st.markdown(assistant_response)

    except Exception as e:
        st.error(f"⚠️ Error: {str(e)}")