import streamlit as st
import openai
import os

# Function to get the API key from Streamlit secrets
def get_api_key():
    try:
        return st.secrets["API_KEY"]
    except KeyError:
        st.error("API_KEY not found in Streamlit secrets. Please add it.")
        return None

# Function to interact with the OpenAI API
def generate_response(prompt, model_name, api_key):
    openai.api_key = api_key
    try:
        completion = openai.ChatCompletion.create(
            model=model_name,
            messages=[{"role": "user", "content": prompt}]
        )
        return completion.choices[0].message.content
    except openai.APIError as e:  # General API Error
        st.error(f"OpenAI API Error with {model_name}: {e}")
        return None
    except openai.RateLimitError as e: # Rate Limit Error
        st.error(f"OpenAI Rate Limit Error with {model_name}: {e}")
        return None
    except openai.AuthenticationError as e: # Authentication Error
        st.error(f"OpenAI Authentication Error with {model_name}: {e}")
        return None
    except Exception as e:  # Catch any other exception
        st.error(f"An unexpected error occurred with {model_name}: {e}")
        return None

# Main Streamlit app
def main():
    st.title("Chatbot with Model Switching")

    # Initialize conversation history in session state
    if "messages" not in st.session_state:
        st.session_state.messages = []

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

    # Get user input
    prompt = st.chat_input("Say something")

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

        # Define model priority
        models = ["gpt-4", "gpt-3.5-turbo"]  # Add more models as needed

        # Get API key
        api_key = get_api_key()
        if not api_key:
            return

        response = None
        for model in models:
            response = generate_response(prompt, model, api_key)
            if response:
                break # If a response is generated, break out of loop.
        
        if response:
            # Add bot message to state
            st.session_state.messages.append({"role": "assistant", "content": response})
            with st.chat_message("assistant"):
                st.markdown(response)

if __name__ == "__main__":
    main()