import streamlit as st
from openai import OpenAI  # OpenAI compatibility
import json
# reference: 
# - Use OpenAI to connect Ollama: https://ollama.com/blog/openai-compatibility
# - Build Chatbot with streamlit: https://streamlit.io/generative-ai
# - Ollama docker: https://hub.docker.com/r/ollama/ollama
# - [TBD] Finetune: https://docs.loopin.network/tutorials/LLM/llama3-finetune


# Clear chat history
def clear_chat():
    st.session_state.messages = []
    st.toast("Chat Cleaned", icon="🧹")

def buffbot():
    # Set up the Streamlit app
    st.markdown("<h1 style='text-align: center; color: #451002;'>BuffBot🦬</h1>", unsafe_allow_html=True)
    st.markdown("<h5 style='text-align: center;'> Your friendly AI chatbot powered by LLM! 🤖 </h3>", unsafe_allow_html=True)
    # Display info and source code
    with st.expander("See Source Code"):
        with open(__file__, "r", encoding="utf-8") as f:
            st.code(f.read(), language="python")
    st.divider()
    # Select AI model for chatbot
    model_options = ["gemini-2.5-flash-lite", "deepseek-chat",  ]
    # on_change callback to clear chat history when model is changed
    selected_model = st.selectbox("**👉Please select a model to start**", model_options, on_change=clear_chat)

    # Initialize session state to store chat history and message count
    if "messages" not in st.session_state:
        st.session_state.messages = []
    # Initialize message count
    if "message_count" not in st.session_state:
        st.session_state.message_count = 0

    # Load API credentials from config.json
    # the config file contains the API key and base URL for the selected model
    """
    {
        "deepseek":{
            "api_url": "https://api.deepseek.com",
            "api_key": "YOUR_API_KEY",
            "model":"deepseek-chat"
        },
        
        "gemini-2.0-flash":{
            "api_key": "YOUR_API_KEY",
            "model": "gemini-2.0-flash"
    },
        
    }  
    """
    # The API key and base URL are loaded based on the selected model
    with open('app_config.json') as config_file:
        config = json.load(config_file)
        
        if selected_model == "gemini-2.5-flash-lite":
            api_base_url = config[selected_model]["api_url"]
            api_key = config[selected_model]["api_key"]
            st.info("Powered by the online [Gemini](https://deepmind.google/technologies/gemini/) API!\
                Just a heads up, you have 10 messages to use.")
            # Set the maximum number of user messages
            MAX_USER_MESSAGES =  10
            
        # deepseek-chat model, online API
        if selected_model == "deepseek-chat":
            api_base_url = config[selected_model]["api_url"]
            api_key = config[selected_model]["api_key"]
            st.info("Powered by the online [DeepSeek](https://www.deepseek.com/) API!\
                Just a heads up, you have 10 messages to use.")
            # Set the maximum number of user messages
            MAX_USER_MESSAGES =  10
            
        # llama3.2:1b model, local API
        if selected_model == "llama3.2:1b":
            api_base_url = config[selected_model]["api_url"]
            api_key = config[selected_model]["api_key"]
            st.info("Powered by local llama3.2:1b model via [Ollama](https://ollama.com/library/llama3.2:1b)!\
                Just a heads up, you have 100 messages to use.")
            MAX_USER_MESSAGES = 100     
        
        if selected_model == "deepseek-r1:1.5b":
            api_base_url = config[selected_model]["api_url"]
            api_key = config[selected_model]["api_key"]
            st.info("Powered by local deepseek-r1:1.5b model via [Ollama](https://ollama.com/library/deepseek-r1:1.5b)!\
                Just a heads up, you have 100 messages to use.")
            MAX_USER_MESSAGES = 100     

    # Initialize OpenAI client to connect with the selected model API
    client = OpenAI(api_key=api_key, base_url=api_base_url)

    # print welcome message
    with st.chat_message("assistant", avatar="🦬"):
        st.markdown("Welcome to BuffBot! What Can I Do for You Today?🌞")
        
    # Display chat history with different avatars for user and AI assistant
    for message in st.session_state.messages:
        if message["role"] == "user":
            avatar="🤠"
        else:
            avatar="🦬"
        with st.chat_message(message["role"], avatar=avatar):
            st.markdown(message["content"])

    if st.session_state.message_count < MAX_USER_MESSAGES:
        # Get user input
        if prompt := st.chat_input("Type your message here..."):
            # Add user message to chat history
            st.session_state.messages.append({"role": "user", "content": prompt})
            st.session_state.message_count += 1
            # Display user message with cowboy avatar
            with st.chat_message("user", avatar="🤠"):
                st.markdown(prompt)

            # Generate reply 
            with st.chat_message("assistant", avatar="🦬"):
                with st.spinner('Thinking...'):
                    # Call the selected model API to generate a response
                    stream = client.chat.completions.create(
                        model=selected_model,
                        messages=[
                            {"role": m["role"], "content": m["content"]}
                            for m in st.session_state.messages
                        ],
                        stream=True, # stream the response
                    )
                # Display the response from the model API
                response = st.write_stream(stream)
            # Add the AI assistant response to the chat history
            st.session_state.messages.append({"role": "assistant", "content": response})

    else:
        st.warning("You have reached the maximum number of messages allowed.\
            Please switch to another model to continue chatting.")
        
    # Clear chat history
    if st.button("Clear Chat"):
        clear_chat()
        st.rerun()