import os
import sys
import streamlit as st
from utils import retrieve_doc
from utils import feedback_doc
from annotated_text import annotated_text
import pandas as pd
# streamlit does not support any states out of the box. On every button click, streamlit reload the whole page 
# and every value gets lost. To keep track of our feedback state we use the official streamlit gist mentioned 
# here https://gist.github.com/tvst/036da038ab3e999a64497f42de966a92
import SessionState

def annotate_answer(answer, context):
    start_idx = context.find(answer)
    end_idx = start_idx+len(answer)
    annotated_text(context[:start_idx],(answer,"ANSWER","#8ef"),context[end_idx:])

def random_questions(df):
     random_row = df.sample(1)
     random_question = random_row["Question Text"].values[0]
     random_answer = random_row["Answer"].values[0]
     return random_question, random_answer

# Define state
state_question = SessionState.get(random_question='Who is the father of Arya Starck?', random_answer='', next_question='false', run_query='false')

# Initalize variables
eval_mode = False
random_question = "Who is the father of Arya Starck?"
eval_labels = os.getenv("EVAL_FILE", "eval_labels_example.csv")

# UI search bar and sidebar      
st.write("# Haystack Demo")
st.sidebar.header("Options")
top_k_reader = st.sidebar.slider("Max. number of answers",min_value=1,max_value=10,value=3,step=1)
top_k_retriever = st.sidebar.slider("Max. number of documents from retriever",min_value=1,max_value=10,value=3,step=1)
eval_mode = st.sidebar.checkbox("Evalution mode")
debug = st.sidebar.checkbox("Show debug info")

# load csv into pandas dataframe
if eval_mode:
    try:
        df = pd.read_csv(eval_labels, sep=";")
    except Exception:
        sys.exit('The eval file was not found. Please check the README for more information.')
    if state_question and hasattr(state_question, 'next_question') and hasattr(state_question, 'random_question') and state_question.next_question:
        random_question = state_question.random_question
        random_answer = state_question.random_answer
    else:
        random_question, random_answer = random_questions(df)
        state_question.random_question = random_question
        state_question.random_answer = random_answer

# Generate new random question
if eval_mode:
    next_question = st.button("Load new question")
    if next_question:
       random_question, random_answer = random_questions(df)
       state_question.random_question = random_question
       state_question.random_answer = random_answer
       state_question.next_question = "true"
       state_question.run_query = "false"
    else:
       state_question.next_question = "false"

# Search bar
question = st.text_input("Please provide your query:",value=random_question)
if state_question and state_question.run_query:
    run_query = state_question.run_query
    st.button("Run")
else:
    run_query = st.button("Run")
    state_question.run_query = run_query

raw_json_feedback = ""

# Get results for query
if run_query:
    with st.spinner("Performing neural search on documents... 🧠 \n "
                    "Do you want to optimize speed or accuracy? \n"
                    "Check out the docs: https://haystack.deepset.ai/docs/latest/optimizationmd "):
        results,raw_json = retrieve_doc(question,top_k_reader=top_k_reader,top_k_retriever=top_k_retriever)

    # Show if we use a question of the given set
    if question == random_question and eval_mode:
        st.write("## Correct answers:")
        random_answer
    
    st.write("## Retrieved answers:")

    # Make every button key unique
    count = 0

    for result in results:
        annotate_answer(result['answer'],result['context'])
        '**Relevance:** ', result['relevance'] , '**Source:** ' , result['source']
        if eval_mode:
            # Define columns for buttons
            button_col1, button_col2, button_col3, button_col4 = st.beta_columns([1,1,1,6])
            if button_col1.button("👍", key=(result['answer'] + str(count)), help="Correct answer"):
                raw_json_feedback = feedback_doc(question,"true",result['document_id'],1,"true",result['answer'],result['offset_start_in_doc'])
                st.success('Thanks for your feedback')
            if button_col2.button("👎", key=(result['answer'] + str(count)), help="Wrong answer and wrong passage"):
                raw_json_feedback = feedback_doc(question,"false",result['document_id'],1,"false",result['answer'],result['offset_start_in_doc'])
                st.success('Thanks for your feedback!')
            if button_col3.button("👎👍", key=(result['answer'] + str(count)), help="Wrong answer, but correct passage"):
                raw_json_feedback = feedback_doc(question,"false",result['document_id'],1,"true",result['answer'],result['offset_start_in_doc'])
                st.success('Thanks for your feedback!')
            count+=1
        st.write("___")
    if debug:
        st.subheader("REST API JSON response")
        st.write(raw_json)