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

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

    if 'cfg' not in st.session_state:
        cfg = OmegaConf.create({
            'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
            'corpus_id': str(os.environ['VECTARA_CORPUS_ID']),
            'api_key': str(os.environ['VECTARA_API_KEY']),
            'prompt_name': 'vectara-experimental-summary-ext-2023-12-11-large',
            'topic': 'Standardized testing in education',
            'human_role': 'in opposition to',
            'bot_role': 'in support of'
        })
        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="Debate Bot", layout="wide")


    # left side content
    with st.sidebar:
        st.markdown(f"## Welcome to Debate Bot.\n\n\n")

        role_options = ['in opposition to', 'in support of']
        cfg.human_role = st.selectbox('Your are:', role_options)
        cfg.bot_role = role_options[1] if cfg.human_role == role_options[0] else role_options[0]

        st.markdown(f"{cfg.topic}.\n\n")
        st.markdown("---")
        st.markdown(
            "## How this works?\n"
            "This app was built with [Vectara](https://vectara.com).\n"
        )
        st.markdown("---")

    if "messages" not in st.session_state.keys():
        st.session_state.messages = [{"role": "assistant", "content": f"Please make your opening statment."}]

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

    # User-provided prompt
    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            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"):
            stream = generate_response(prompt, cfg.bot_role, cfg.topic) 
            response = st.write_stream(stream) 
            message = {"role": "assistant", "content": response}
            st.session_state.messages.append(message)
    
if __name__ == "__main__":
    launch_bot()