import os  # Make sure this is included
import streamlit as st
from streamlit_chat import message
from openai import OpenAI

# Load API key securely
API_KEY = os.getenv("NV_API_KEY", "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")
if not API_KEY:
    st.error("API key is missing! Please set NV_API_KEY as an environment variable.")
    st.stop()

client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=API_KEY
)

# Initialize the conversation history
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "Hi! I'm your chatbot. How can I assist you today?"}]

# Function to get a response from OpenAI
def get_chatbot_response(user_input):
    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",  # Use GPT-4 if you have access for better responses
            messages=st.session_state.messages + [{"role": "user", "content": user_input}],
        )
        return response['choices'][0]['message']['content']
    except Exception as e:
        return f"Error: {str(e)}"

# App layout
st.set_page_config(page_title="ChatGPT-like Chatbot", page_icon="🤖", layout="wide")
st.title("ChatGPT-like Chatbot 🤖")

# Display the chat conversation
for msg in st.session_state.messages:
    if msg["role"] == "assistant":
        message(msg["content"], is_user=False, key=f"assistant_{msg['content']}")
    else:
        message(msg["content"], is_user=True, key=f"user_{msg['content']}")

# User input section
user_input = st.text_input(
    "Type your message:", key="user_input", placeholder="Type here and press Enter..."
)

if user_input:
    # Append user message to session state
    st.session_state.messages.append({"role": "user", "content": user_input})
    
    # Get chatbot response
    chatbot_response = get_chatbot_response(user_input)
    
    # Append chatbot response to session state
    st.session_state.messages.append({"role": "assistant", "content": chatbot_response})
    
    # Clear the input field
    st.session_state.user_input = ""