import streamlit as st
from ir import SearchEngine
from qa import QAEngine


# initialize a search tool
if not 'ir' in st.session_state:
	st.session_state.ir = SearchEngine()
	st.session_state.qa = QAEngine()
st.title('Q&A Engine')

# render indexed prompts
st.write('---')
st.subheader('Indexed Context:')
if len(st.session_state.ir.prompts) > 0: # if there are any prompts, proceed
	for id, prompt in enumerate(st.session_state.ir.prompts, 1):
		st.info(f'{id}. {prompt}')
else: # if there are no prompts
	st.write('No prompts yet.')


# render inputable prompt
st.write('---')
st.subheader('The Question:')
input_prompt = st.text_area('Enter a prompt to save or match:', height=150)


# initialize two buttons
col1, col2, col3 = st.columns(3)

with col1:
	if st.button('Save', use_container_width=True): # it should add a new prompt to the ir
		st.session_state.ir.add(input_prompt)
		st.success('Saved successfully!')
		st.rerun()

with col2:
	if st.button('Answer', use_container_width=True): # it should match a new prompt against others and create an answer based on a context
		if len(st.session_state.ir.vectors) > 0:
			match = st.session_state.ir.search(input_prompt)
			answer = st.session_state.qa.answer(question=input_prompt, context=match)
			st.success(f'Result: {answer}')
		else:
			st.warning('Nothing to match against.')

with col3:
	if st.button('Remove', use_container_width=True): # it should remove a specified prompt
		try:
			idx = int(input_prompt)
			st.session_state.ir.prompts.pop(idx - 1)
			st.session_state.ir.vectors.pop(idx - 1)
		except:
			pass

		st.rerun()