import sys
import toml
from omegaconf import OmegaConf
from query import VectaraQuery
import os

import streamlit as st
from PIL import Image

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

def process_prompt(prompt):
    # Your logic to process the prompt and generate response
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            response = generate_response(prompt) 
            st.write(response) 
    message = {"role": "assistant", "content": response}
    st.session_state.messages.append(message)

def launch_bot():

    if 'cfg' not in st.session_state:
        corpus_ids = str(os.environ['corpus_ids']).split(',')
        questions = list(eval(os.environ['examples']))
        cfg = OmegaConf.create({
            'customer_id': str(os.environ['customer_id']),
            'corpus_ids': corpus_ids,
            'api_key': str(os.environ['api_key']),
            'title': os.environ['title'],
            'description': os.environ['description'],
            'examples': questions,
            'source_data_desc': os.environ['source_data_desc']
        })
        st.session_state.cfg = cfg
        st.session_state.vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_ids)

    cfg = st.session_state.cfg
    st.set_page_config(page_title=cfg.title, layout="wide")

    # left side content
    with st.sidebar:
        image = Image.open('Vectara-logo.png')
        st.markdown(f"## Welcome to {cfg.title}\n\n"
                    f"With this demo uses Retrieval Augmented Generation to ask questions about {cfg.source_data_desc}\n\n")

        st.markdown("---")
        st.markdown(
            "## How this works?\n"
            "This app was built with [Vectara](https://vectara.com).\n"
            "Vectara's [Indexing API](https://docs.vectara.com/docs/api-reference/indexing-apis/indexing) was used to ingest the data into a Vectara corpus (or index).\n\n"
            "This app uses Vectara [Chat API](https://docs.vectara.com/docs/console-ui/vectara-chat-overview) to query the corpus and present the results to you, answering your question.\n\n"
        )
        st.markdown("---")
        st.image(image, width=250)

    st.markdown(f"<center> <h2> Vectara chat demo: {cfg.title} </h2> </center>", unsafe_allow_html=True)
    st.markdown(f"<center> <h4> {cfg.description} <h4> </center>", unsafe_allow_html=True)

    if "messages" not in st.session_state.keys():
        st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
    
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])

     # Always allow typing in the chat input box
    user_input = st.text_input("Type your question here...", key="user_input")

    # Display example questions only during the first round
    if len(st.session_state.messages) == 0:  # Only show examples in the first round
        for example in cfg.examples:
            if st.button(example):
                user_input = example  # Directly use the example as input
                # Process the input immediately
                response = generate_response(user_input)
                st.session_state.messages.append(user_input)
                st.session_state.messages.append(response)
                st.experimental_rerun()  # Rerun to refresh and show the response

    if st.button("Submit"):
        if user_input:  # Ensure there's something to process
            response = generate_response(user_input)
            st.session_state.messages.append(user_input)
            st.session_state.messages.append(response)
            # Clear the input (workaround as direct clearing isn't supported)
            st.experimental_rerun()  # Rerun to refresh and clear the input box
            
if __name__ == "__main__":
    launch_bot()