# Run Streamlit app
import streamlit as st
import os
from dotenv import load_dotenv
from groq import Groq

# Load environment variables
load_dotenv()

# Initialize Groq client
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
    st.error("API key not found in the environment. Please check your .env file.")
    st.stop()

client = Groq(api_key=api_key)

# Streamlit app layout
st.title("Chat with Groq")
st.write("Ask anything, and the Groq model will respond!")

# Sidebar for settings
st.sidebar.title("Settings")
model_name = st.sidebar.text_input("Model Name", value="llama3-8b-8192")

# Main chat interface
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "system", "content": "You are a helpful assistant."}
    ]

for message in st.session_state.messages:
    if message["role"] == "user":
        st.markdown(f"**User:** {message['content']}")
    elif message["role"] == "assistant":
        st.markdown(f"**Assistant:** {message['content']}")

# Input box for user query
user_input = st.text_input("Your Message", key="input_box")

# Handle user input and get response
if st.button("Send") and user_input.strip():
    # Add user message to session
    st.session_state.messages.append({"role": "user", "content": user_input})

    # Call Groq API
    try:
        chat_completion = client.chat.completions.create(
            messages=st.session_state.messages,
            model=model_name,
        )
        assistant_response = chat_completion.choices[0].message.content

        # Add assistant message to session
        st.session_state.messages.append({"role": "assistant", "content": assistant_response})

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

    # Clear input box
    st.session_state.input_box = ""

# Button to clear chat history
if st.sidebar.button("Clear Chat History"):
    st.session_state.messages = [
        {"role": "system", "content": "You are a helpful assistant."}
    ]
    st.success("Chat history cleared.")