import streamlit as st
import uuid
import os
import re
import sys
sys.path.insert(1, "/".join(os.path.realpath(__file__).split("/")[0:-2])+"/semantic_search")
sys.path.insert(1, "/".join(os.path.realpath(__file__).split("/")[0:-2])+"/RAG")
sys.path.insert(1, "/".join(os.path.realpath(__file__).split("/")[0:-2])+"/utilities")
import boto3
import requests
from boto3 import Session
import botocore.session
import json
import random
import string
import rag_DocumentLoader 
import rag_DocumentSearcher
import pandas as pd
from PIL import Image 
import shutil
import base64
import time
import botocore
from requests_aws4auth import AWS4Auth
import colpali
from requests.auth import HTTPBasicAuth



st.set_page_config(
    #page_title="Semantic Search using OpenSearch",
    layout="wide",
    page_icon="images/opensearch_mark_default.png"
)
parent_dirname = "/".join((os.path.dirname(__file__)).split("/")[0:-1])
USER_ICON = "images/user.png"
AI_ICON = "images/opensearch-twitter-card.png"
REGENERATE_ICON = "images/regenerate.png"
s3_bucket_ = "pdf-repo-uploads"
           
# polly_client = boto3.Session(
#             region_name='us-east-1').client('polly')

# Check if the user ID is already stored in the session state
if 'user_id' in st.session_state:
    user_id = st.session_state['user_id']
    

# If the user ID is not yet stored in the session state, generate a random UUID
else:
    user_id = str(uuid.uuid4())
    st.session_state['user_id'] = user_id


if 'session_id' not in st.session_state:
    st.session_state['session_id'] = ""
    
if "chats" not in st.session_state:
    st.session_state.chats = [
        {
            'id': 0,
            'question': '',
            'answer': ''
        }
    ]

if "questions_" not in st.session_state:
    st.session_state.questions_ = []
    
    
if "show_columns" not in st.session_state:
    st.session_state.show_columns = False
    
if "answer_ready" not in st.session_state:
    st.session_state.answer_ready = False

if "answers_" not in st.session_state:
    st.session_state.answers_ = []

if "input_index" not in st.session_state:
    st.session_state.input_index = "globalwarming"#"hpijan2024hometrack"#"hpijan2024hometrack"#"covid19ie"#"choosetheknnalgorithmforyourbillionscaleusecasewithopensearchawsbigdatablog"
    
if "input_is_rerank" not in st.session_state:
    st.session_state.input_is_rerank = True
    
if "input_is_colpali" not in st.session_state:
    st.session_state.input_is_colpali = False

if "input_copali_rerank" not in st.session_state:
    st.session_state.input_copali_rerank = False
    
if "input_table_with_sql" not in st.session_state:
    st.session_state.input_table_with_sql = False
    
if "input_query" not in st.session_state:
    if(st.session_state.input_index == "globalwarming"):
        st.session_state.input_query="What is the projected energy percentage from renewable sources in future ?"#"What is the projected energy percentage from renewable sources in future?"#"Which city in United Kingdom has the highest average housing price ?"#"How many aged above 85 years died due to covid ?"# What is the projected energy from renewable sources ?"
    if(st.session_state.input_index == "hpijan2024hometrack"):
        st.session_state.input_query="which city has the highest average housing price in UK ?"#"What is the projected energy percentage from renewable sources in future?"#"Which city in United Kingdom has the highest average housing price ?"#"How many aged above 85 years died due to covid ?"# What is the projected energy from renewable sources ?"
    if(st.session_state.input_index == "covid19ie"):
        st.session_state.input_query="How many aged above 85 years died due to covid ?"#"What is the projected energy percentage from renewable sources in future?"#"Which city in United Kingdom has the highest average housing price ?"#"How many aged above 85 years died due to covid ?"# What is the projected energy from renewable sources ?"


st.markdown("""
    <style>
    [data-testid=column]:nth-of-type(2) [data-testid=stVerticalBlock]{
        gap: 0rem;
    }
    [data-testid=column]:nth-of-type(1) [data-testid=stVerticalBlock]{
        gap: 0rem;
    }
    </style>
    """,unsafe_allow_html=True)

################# using boto3 credentials ###################

awsauth = HTTPBasicAuth('master',st.secrets['ml_search_demo_api_access'])




def write_logo():
    col1, col2, col3 = st.columns([5, 1, 5])
    with col2:
        st.image(AI_ICON, use_column_width='always') 

def write_top_bar():
    col1, col2 = st.columns([77,23])
    with col1:
        st.write("")
        st.header("Chat with your data",divider='rainbow')
        
        #st.image(AI_ICON, use_column_width='always')
    
    with col2:
        st.write("")
        st.write("")
        clear = st.button("Clear")
    st.write("")
    st.write("")
    return clear

clear = write_top_bar()

if clear:
    st.session_state.questions_ = []
    st.session_state.answers_ = []
    st.session_state.input_query=""



def handle_input(state,dummy):
    if(state == 'colpali_show_similarity_map'):
        st.session_state.show_columns = True    
    print("Question: "+st.session_state.input_query)
    print("-"*20)
    print("\n\n")
    # if(st.session_state.input_query==''):
    #     return ""
    inputs = {}
    for key in st.session_state:
        if key.startswith('input_'):
            inputs[key.removeprefix('input_')] = st.session_state[key]
    st.session_state.inputs_ = inputs
    question_with_id = {
        'question': inputs["query"],
        'id': len(st.session_state.questions_)
    }
    st.session_state.questions_.append(question_with_id)
    if(st.session_state.input_is_colpali):
        out_ = colpali.colpali_search_rerank(st.session_state.input_query)
        
    else:
        out_ = rag_DocumentSearcher.query_(awsauth, inputs, st.session_state['session_id'],st.session_state.input_rag_searchType)
    st.session_state.answers_.append({
        'answer': out_['text'],
        'source':out_['source'],
        'id': len(st.session_state.questions_),
        'image': out_['image'],
        'table':out_['table']
    })
    st.session_state.input_query=""


def write_user_message(md):
    col1, col2 = st.columns([3,97])
    
    with col1:
        st.image(USER_ICON, use_column_width='always')
    with col2:
        st.markdown("<div style='color:#e28743';font-size:18px;padding:3px 7px 3px 7px;borderWidth: 0px;borderColor: red;borderStyle: solid;width: fit-content;height: fit-content;border-radius: 10px;font-style: italic;'>"+md['question']+"</div>", unsafe_allow_html = True)
       


def render_answer(question,answer,index,res_img):
    
    
    col1, col2, col_3 = st.columns([4,74,22])
    with col1:
        st.image(AI_ICON, use_column_width='always')
    with col2:
        ans_ = answer['answer']
        st.write(ans_)
        rdn_key_1 = ''.join([random.choice(string.ascii_letters)
                              for _ in range(10)])

            
        if(st.session_state.input_is_colpali):
            placeholder__ = st.empty()
            placeholder__.button("Show similarity map",key=rdn_key_1,on_click=handle_input,args=('colpali_show_similarity_map',True))
   

    colu1,colu2,colu3 = st.columns([4,82,20])
    with colu2:
        with st.expander("Relevant Sources:"):
            with st.container():
                if(len(res_img)>0):
                    #with st.expander("Images:"):
                        
                    idx = 0
                    for i in range(0,len(res_img)):
                        
                        if(st.session_state.input_is_colpali):
                            if(len(res_img)>1):
                                cols_per_row = 3
                                st.session_state.image_placeholder=st.empty()
                                with st.session_state.image_placeholder.container():
                                    row = st.columns(cols_per_row)
                                for j, item in enumerate(res_img[i:i+cols_per_row]):
                                    with row[j]:
                                        st.markdown(
                                            """
                                            <div style='max-height:500px; overflow:auto; border:1px solid #444; padding:4px;'>
                                            """,
                                            unsafe_allow_html=True
                                        )
                                        st.image(item['file'])
                                        st.markdown("</div>", unsafe_allow_html=True)
                                
                            else:
                                st.session_state.image_placeholder = st.empty()
                                with st.session_state.image_placeholder.container():
                                    col3_,col4_,col5_ = st.columns([33,33,33])
                                    with col3_:
                                        st.markdown(
                                            """
                                            <div style='max-height:500px; overflow:auto; border:1px solid #444; padding:4px;'>
                                            """,
                                            unsafe_allow_html=True
                                        )
                                        st.image(res_img[i]['file'])
                        else:
                            if(res_img[i]['file'].lower()!='none' and idx < 1):
                                col3,col4,col5 = st.columns([33,33,33])
                                cols = [col3,col4]
                                img = res_img[i]['file'].split(".")[0]
                                caption = res_img[i]['caption']
                            
                                with cols[idx]:
                                    st.markdown(
                                            """
                                            <div style='max-height:500px; overflow:auto; border:1px solid #444; padding:4px;'>
                                            """,
                                            unsafe_allow_html=True
                                        )
                                    st.image(parent_dirname+"/figures/"+st.session_state.input_index+"/"+img+".jpg")
                                    st.markdown("</div>", unsafe_allow_html=True)
                            
                            idx = idx+1
                if(len(answer["table"] )>0):
                    #with st.expander("Table:"):
                    df = pd.read_csv(answer["table"][0]['name'],skipinitialspace = True, on_bad_lines='skip',delimiter='`')
                    df.fillna(method='pad', inplace=True)
                    st.table(df)
                #with st.expander("Raw sources:"):
                st.write(answer["source"])
    with col_3:
        if(index == len(st.session_state.questions_)):

            rdn_key = ''.join([random.choice(string.ascii_letters)
                              for _ in range(10)])
            currentValue = ''.join(st.session_state.input_rag_searchType)+str(st.session_state.input_is_rerank)+str(st.session_state.input_table_with_sql)+st.session_state.input_index
            oldValue = ''.join(st.session_state.inputs_["rag_searchType"])+str(st.session_state.inputs_["is_rerank"])+str(st.session_state.inputs_["table_with_sql"])+str(st.session_state.inputs_["index"])
            def on_button_click():
                if(currentValue!=oldValue or 1==1): 
                    st.session_state.input_query = st.session_state.questions_[-1]["question"]
                    st.session_state.answers_.pop()
                    st.session_state.questions_.pop()
                    handle_input("regenerate_",None)
                    with placeholder.container():
                        render_all()
            
                

            if("currentValue"  in st.session_state):
                del st.session_state["currentValue"]

            try:
                del regenerate
            except:
                pass  

            placeholder__ = st.empty()
            placeholder__.button("🔄",key=rdn_key,on_click=on_button_click)
            
        
     
#Each answer will have context of the question asked in order to associate the provided feedback with the respective question
def write_chat_message(md, q,index):
    res_img = md['image']
    #st.session_state['session_id'] = res['session_id']   to be added in memory
    chat = st.container()
    with chat:
        render_answer(q,md,index,res_img)
    
def render_all():  
    index = 0
    for (q, a) in zip(st.session_state.questions_, st.session_state.answers_):
        index = index +1
        
        write_user_message(q)
        write_chat_message(a, q,index)

placeholder = st.empty()
with placeholder.container():
  render_all()

st.markdown("")
col_2, col_3 = st.columns([75,20])
    
with col_2:
    input = st.text_input( "Ask here",label_visibility = "collapsed",key="input_query")
with col_3:
    #hidden = st.button("RUN",disabled=True,key = "hidden")
    play = st.button("Go",on_click=handle_input,key = "play",args=(None,None),help = 'Enter the questions and click on "GO"')
with st.sidebar:
    st.page_link("app.py", label=":orange[Home]", icon="🏠")
    st.subheader(":blue[Sample Data]")
    coln_1,coln_2 = st.columns([70,30])
    with coln_1:
        index_select = st.radio("Choose one index",["Global Warming stats","UK Housing","Covid19 impacts on Ireland"],key="input_rad_index")
    with coln_2:
        st.markdown("<p style='font-size:15px'>Preview file</p>",unsafe_allow_html=True)
        st.write("[:eyes:](https://github.com/aws-samples/AI-search-with-amazon-opensearch-service/blob/b559f82c07dfcca973f457c0a15d6444752553ab/rag/sample_pdfs/global_warming.pdf)")
        st.write("[:eyes:](https://github.com/aws-samples/AI-search-with-amazon-opensearch-service/blob/b559f82c07dfcca973f457c0a15d6444752553ab/rag/sample_pdfs/HPI-Jan-2024-Hometrack.pdf)")
        st.write("[:eyes:](https://github.com/aws-samples/AI-search-with-amazon-opensearch-service/blob/b559f82c07dfcca973f457c0a15d6444752553ab/rag/sample_pdfs/covid19_ie.pdf)")
    st.markdown("""
    <style>
    [data-testid=column]:nth-of-type(2) [data-testid=stVerticalBlock]{
        gap: 0rem;
    }
    [data-testid=column]:nth-of-type(1) [data-testid=stVerticalBlock]{
        gap: 0rem;
    }
    </style>
    """,unsafe_allow_html=True)   
    with st.expander("Sample questions:"):
        st.markdown("<span style = 'color:#FF9900;'>Global Warming stats</span> - What is the projected energy percentage from renewable sources in future?",unsafe_allow_html=True)
        st.markdown("<span style = 'color:#FF9900;'>UK Housing</span> - which city has the highest average housing price in UK ?",unsafe_allow_html=True)
        st.markdown("<span style = 'color:#FF9900;'>Covid19 impacts</span> - How many aged above 85 years died due to covid ?",unsafe_allow_html=True)
        
    
    
    # st.subheader(":blue[Your multi-modal documents]")
    # pdf_doc_ = st.file_uploader(
    #     "Upload your PDFs here and click on 'Process'", accept_multiple_files=False)
                    
                
    # pdf_docs = [pdf_doc_]
    # if st.button("Process"):
    #     with st.spinner("Processing"):
    #         if os.path.isdir(parent_dirname+"/pdfs") == False:
    #             os.mkdir(parent_dirname+"/pdfs")
            
    #         for pdf_doc in pdf_docs:
    #             print(type(pdf_doc))
    #             pdf_doc_name = (pdf_doc.name).replace(" ","_")
    #             with open(os.path.join(parent_dirname+"/pdfs",pdf_doc_name),"wb") as f: 
    #                 f.write(pdf_doc.getbuffer())  
                    
    #             request_ = { "bucket": s3_bucket_,"key": pdf_doc_name}
    #             # if(st.session_state.input_copali_rerank):
    #             #     copali.process_doc(request_)
    #             # else:
    #             rag_DocumentLoader.load_docs(request_)
    #             print('lambda done')
    #     st.success('you can start searching on your PDF')
        
     ############## haystach demo temporary addition ############       
    #if(pdf_doc_ is None or pdf_doc_ == ""):
    if(index_select == "Global Warming stats"):
        st.session_state.input_index = "globalwarming"
    if(index_select == "Covid19 impacts on Ireland"):
        st.session_state.input_index = "covid19ie"#"choosetheknnalgorithmforyourbillionscaleusecasewithopensearchawsbigdatablog"
    if(index_select == "BEIR"):
        st.session_state.input_index = "2104"
    if(index_select == "UK Housing"):
        st.session_state.input_index = "hpijan2024hometrack"
          
    # custom_index = st.text_input("If uploaded the file already, enter the original file name", value = "")
    # if(custom_index!=""):
    #     st.session_state.input_index = re.sub('[^A-Za-z0-9]+', '', (custom_index.lower().replace(".pdf","").split("/")[-1].split(".")[0]).lower())
    
    
    
    st.subheader(":blue[Retriever]")
    search_type = st.multiselect('Select the Retriever(s)',
    ['Keyword Search',
    'Vector Search', 
    'Sparse Search',
    ],
    ['Vector Search'],

    key = 'input_rag_searchType',
    help = "Select the type of Search, adding more than one search type will activate hybrid search"#\n1. Conversational Search (Recommended) - This will include both the OpenSearch and LLM in the retrieval pipeline \n (note: This will put opensearch response as context to LLM to answer) \n2. OpenSearch vector search - This will put only OpenSearch's vector search in the pipeline, \n(Warning: this will lead to unformatted results )\n3. LLM Text Generation - This will include only LLM in the pipeline, \n(Warning: This will give hallucinated and out of context answers)"
    )
    
    re_rank = st.checkbox('Re-rank results', key = 'input_re_rank', disabled = False, value = True, help = "Checking this box will re-rank the results using a cross-encoder model")
     
    if(re_rank):
        st.session_state.input_is_rerank = True
    else:
        st.session_state.input_is_rerank = False
    
    st.subheader(":blue[Multi-vector retrieval]")
    
    colpali_search_rerank =  st.checkbox('Try Colpali multi-vector retrieval on the [sample dataset](https://huggingface.co/datasets/vespa-engine/gpfg-QA)', key = 'input_colpali', disabled = False, value = False, help = "Checking this box will use colpali as the embedding model and retrieval is performed using multi-vectors followed by re-ranking using MaxSim")
        
    if(colpali_search_rerank):
        st.session_state.input_is_colpali = True
        #st.session_state.input_query = ""
    else:
        st.session_state.input_is_colpali = False

    #st.button("Show similarity map",key="input_sim_map",on_click=handle_input,args=('colpali_show_similarity_map',True))
       
    with st.expander("Sample questions for Colpali retriever:"):
        st.write("1. Proportion of female new hires 2021-2023? \n\n 2. First-half 2021 return on unlisted real estate investments? \n\n 3. Trend of the fund's expected absolute volatility between January 2014 and January 2016? \n\n 4. Fund return percentage in 2017? \n\n 5. Annualized gross return of the fund from 1997 to 2008?")