import streamlit as st
import os
from groq import Groq

# You will set API key directly in Colab cell environment (no .env file needed)
api_key = os.environ['GROQ_API']

if not api_key:
    st.error("Please set the GROQ_API_KEY environment variable.")
    st.stop()

# Initialize GROQ client
client = Groq(api_key=api_key)

# Streamlit app
st.set_page_config(page_title="GROQ Chatbot", page_icon="🤖")
st.title("🤖 GROQ Chatbot")

# Session state to store messages
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display previous messages
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# User input
user_prompt = st.chat_input("Type your message...")

if user_prompt:
    # Add user message to session state
    st.session_state.messages.append({"role": "user", "content": user_prompt})
    
    with st.chat_message("user"):
        st.markdown(user_prompt)

    # Send the conversation to GROQ API
    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {"role": "user", "content": user_prompt}
            ],
            model="llama-3.3-70b-versatile",
        )

        bot_reply = chat_completion.choices[0].message.content

        # Add assistant response to session state
        st.session_state.messages.append({"role": "assistant", "content": bot_reply})

        with st.chat_message("assistant"):
            st.markdown(bot_reply)

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