from omegaconf import OmegaConf
from query import VectaraQuery
import streamlit as st
import os
from PIL import Image

def launch_bot():
    def generate_response(question):
        response = vq.submit_query(question)
        return response

    def reset():
        st.session_state.messages = [{"role": "assistant", "content": "Please ask your question about drink names.", "avatar": '🦖'}]
        st.session_state.vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_id, cfg.prompt_name)

    if 'cfg' not in st.session_state:
        cfg = OmegaConf.create({
            'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
            'corpus_id': '46',  # Fixed corpus ID for drink names
            'api_key': str(os.environ['VECTARA_API_KEY']),
            'prompt_name': 'vectara-experimental-summary-ext-2023-12-11-large',
        })
        st.session_state.cfg = cfg
        st.session_state.vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_id, cfg.prompt_name)

    cfg = st.session_state.cfg
    vq = st.session_state.vq
    st.set_page_config(page_title="Drink Name Query Bot", layout="wide")

    # Left side content
    with st.sidebar:
        image = Image.open('Vectara-logo.png')
        st.image(image, width=250)
        st.markdown(f"## Welcome to Drink Name Query Bot.\n\n\n")

        if st.button('Start Over'):
            reset()

        st.markdown("---")
        st.markdown(
            "## How this works?\n"
            "This app was built with [Vectara](https://vectara.com).\n\n"
            "It demonstrates the use of the Chat functionality along with custom prompts and GPT4-Turbo (as part of our [Scale plan](https://vectara.com/pricing/))"
        )
        st.markdown("---")

    if "messages" not in st.session_state.keys():
        reset()

    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"], avatar=message["avatar"]):
            st.write(message["content"])

    # User-provided prompt
    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt, "avatar": '🧑‍💻'})
        with st.chat_message("user", avatar='🧑‍💻'):
            st.write(prompt)

    # Generate a new response if last message is not from assistant
    if st.session_state.messages[-1]["role"] != "assistant":
        with st.chat_message("assistant", avatar='🤖'):
            response = generate_response(prompt) 
            st.write(response)
            message = {"role": "assistant", "content": response, "avatar": '🤖'}
            st.session_state.messages.append(message)
    
if __name__ == "__main__":
    launch_bot()