import streamlit as st
from openai import OpenAI
import os
from dotenv import load_dotenv

load_dotenv()

# Set up OpenAI API key
api_key = os.getenv("OPENAI_API_KEY")  # Replace with your API key if not using an environment variable
client = OpenAI(api_key=api_key)

# Function to query OpenAI
def query_openai(prompt, model="gpt-4o-mini"):
    try:
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Respond in 20 words."},
                {"role": "user", "content": prompt}
            ]
        )
        return response.choices[0].message['content']
    except Exception as e:
        return f"Error: {e}"

# Streamlit app
st.title("Simple OpenAI Chatbot")

# Initialize the chat history
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": "You are a helpful assistant. Respond in 20 words"}]

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

# User input
user_input = st.text_input("Type your message:")

# Handle user input
if st.button("Send"):
    if user_input.strip():
        # Add user input to chat history
        st.session_state.messages.append({"role": "user", "content": user_input})

        # Get bot response
        bot_response = query_openai(user_input)

        # Add bot response to chat history
        st.session_state.messages.append({"role": "assistant", "content": bot_response})

# Clear chat
if st.button("Clear Chat"):
    st.session_state.messages = [{"role": "system", "content": "You are a helpful assistant. Respond in 20 words"}]