import os
import streamlit as st
from groq import Groq

# Get Groq API key from environment variable (set as HF secret)
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
client = Groq(api_key=GROQ_API_KEY)

# Streamlit UI
st.set_page_config(page_title="Groq AI Chatbot", page_icon="🤖")
st.title("🤖 Groq AI Chatbot (Streamlit)")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state["messages"] = []

# Display chat history
for msg in st.session_state["messages"]:
    if msg["role"] == "user":
        st.markdown(f"**You:** {msg['content']}")
    else:
        st.markdown(f"**AI:** {msg['content']}")

# User input
prompt = st.text_input("Type your question here:")

if prompt:
    st.session_state["messages"].append({"role": "user", "content": prompt})
    try:
        completion = client.chat.completions.create(
            model="llama-3.1-8b-instant",  # supported Groq model
            messages=st.session_state["messages"]
        )
        reply = completion.choices[0].message.content
    except Exception as e:
        reply = f"(Error: {e})"

    st.session_state["messages"].append({"role": "assistant", "content": reply})
    st.experimental_rerun()