import streamlit as st
import logging
import os

from annotated_text import annotation
from json import JSONDecodeError
from markdown import markdown
from utils.config import parser
from utils.haystack import start_document_store, start_haystack_extractive, start_haystack_rag, query
from utils.ui import reset_results, set_initial_state

try:
    args = parser.parse_args()
    document_store = start_document_store(type = args.store)
    if args.task == 'extractive':
        pipeline = start_haystack_extractive(document_store)
    else:
        pipeline = start_haystack_rag(document_store)

    set_initial_state()

    st.write('# '+args.name)

    # Search bar
    question = st.text_input("Ask a question", value=st.session_state.question, max_chars=100, on_change=reset_results)
    #question = "what is Pi?"
    
    run_pressed = st.button("Run")
    #run_pressed = True

    run_query = (
        run_pressed or question != st.session_state.question
    )

    # Get results for query
    if run_query and question:
        reset_results()
        st.session_state.question = question
        with st.spinner("🔎    Running your pipeline"):
            try:
                st.session_state.results = query(pipeline, question)
            except JSONDecodeError as je:
                st.error(
                    "👓    An error occurred reading the results. Is the document store working?"
                )    
            except Exception as e:
                logging.exception(e)
                st.error("🐞    An error occurred during the request.")
            
                

    if st.session_state.results:
        results = st.session_state.results

        if args.task == 'extractive':
            answers = results['answers']
            for count, answer in enumerate(answers):
                if answer.answer:
                    text, context = answer.answer, answer.context
                    start_idx = context.find(text)
                    end_idx = start_idx + len(text)
                    st.write(
                        f" Answer: {markdown(context[:start_idx] + str(annotation(body=text, label='ANSWER', background='#964448', color='#ffffff')) + context[end_idx:])}",
                        unsafe_allow_html=True,
                    )
                else:
                    st.info(
                        "🤔    Haystack is unsure whether any of the documents contain an answer to your question. Try to reformulate it!"
                    )
        elif args.task == 'rag':
            st.write(f" Answer: {results['results'][0]}")
        
                # Extract and display information from the 'documents' list
        retrieved_documents = results['documents']
        st.subheader("Retriever Results:")
        for document in retrieved_documents:
            st.write(f"Document Name: {document.meta['name']}")
            st.write(f"Score: {document.score}")
            st.write(f"Text: {document.content}")
except SystemExit as e:
    # This exception will be raised if --help or invalid command line arguments
    # are used. Currently streamlit prevents the program from exiting normally
    # so we have to do a hard exit.
    os._exit(e.code)