from sentence_transformers import SentenceTransformer
from utils_st import load_data,clean_whitespace
import streamlit as st


#embed_model, reranker = load_models()
hybrid_retriever = load_data()


def stream(reranked_nodes,text_size=700):
    nodes_dict={}
    for nod in reranked_nodes:
        file_name = nod.metadata["file_name"]
        if file_name not in nodes_dict:  nodes_dict[file_name]=[]
        if nod not in nodes_dict[file_name]:
            nodes_dict[file_name].append(nod)
    for rank,i_di in enumerate(nodes_dict):
        i_di=nodes_dict[i_di]
        title = i_di[0].metadata['title']
        file_name = i_di[0].metadata['file_name']
        summary = i_di[0].metadata['text']
        url = i_di[0].metadata['doc_url'] 
        st.write(f"**Rank {rank+1}:** {file_name} ")
        st.write(f"- Tittle: [{title}](%s)"% url)
        #st.write("check out this [link](%s)" % url)
        with st.expander(f"Summary"):
            st.write(f"{summary}")
        with st.expander(f"Extra Text(s) "):
            for n_extra,t in enumerate(i_di[:5]):
                page_n=t.metadata['page_label'] if "page_label" in t.metadata else 'Unknown'
                st.write(f"- **Found in page n°{page_n}** - Extra text **n°{n_extra+1}:**")
                st.write(f"\t     {clean_whitespace(t.text[:text_size])}...")
                st.markdown("""---""")
        st.markdown("""---""")



# Function to perform search and return sorted documents
def perform_search(query):
    if query:
        retrieved_nodes = hybrid_retriever.retrieve(query)
        
        return retrieved_nodes
    else:
        return []

# Main Streamlit app
def main():
    #st.sidebar.title("Pagination")
    #st.sidebar.write( f"Totale Pages :{total_pages}")
   
    st.title("EIGA Search IMS Business")
    query = st.text_input("Enter your question:")

    if st.button("Search") :
        sorted_docs = perform_search(query)
        st.session_state.sorted_docs = sorted_docs
        
        
    else:
        sorted_docs = st.session_state.get("sorted_docs", [])

    if sorted_docs:
        stream(sorted_docs,700)
        query = st.write("Powered By: Arthur SATOUF, Air liquide R&D")
        #st.write(f"Current Page Number: {page_number}")
        

    

# Run the app
if __name__ == "__main__":
    main()