import gradio as gr
from embed_with_db import get_all_collections, VECTORDB_STORE, config
from vectorize import VectorDataBase

def respond(message, chat_history, collection_name):
    chain = VECTORDB_STORE(collection_name).chain()
    res = chain.invoke(message)
    chat_history.append((message, res))
    return "", chat_history


def embed_and_store(password, collection_name, file_type, file_fields, context,page_start):
    if password == config['PASSWORD_DB']:
        if str(file_type)== 'string':
            file_fields = context
        vector_db = VectorDataBase(file_fields, collection_name, file_type, page_start=page_start)
        vector_db.embedding_with_loop()
        return file_fields,context
    else:
        raise Exception('Something went wrong')
def update_interface(file_type):
    if file_type == 'PDF' or file_type == 'TEXT':
        return gr.Textbox(visible= False),gr.File(label = 'Select the file',interactive= True,visible= True)
    else:
        return gr.Textbox(visible = True, label= 'Enter the Context', interactive= True),gr.File(visible= False)

with gr.Blocks() as demo:
    with gr.Tab('Personal Chat bot'):
        gr.Markdown("""
    <div align='center'>RAG Application with Open Source models</div>
       <div align='center'>
    > You could ask anything about Me & Data Science. I hope it will find you well
        </div>
            """)
        db_collection = gr.Dropdown(
                list(get_all_collections().values()), label="Select Collection for the retriever",
                  value= 'Data scientist',
                  allow_custom_value=True)
        chatbot = gr.Chatbot(height=480)  # Just to fit the notebook
        msg = gr.Textbox(label="Prompt", interactive= True)
        btn = gr.Button("Submit")
        clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")
        btn.click(respond, inputs=[msg, chatbot, db_collection], outputs=[msg, chatbot])
        msg.submit(respond, inputs=[msg, chatbot,db_collection], outputs=[msg, chatbot])  # Press enter to submit

    with gr.Tab('Data Base and Embedding Store'):
        gr.Markdown("""
    <div align='center'>Store the Document | String in Database</div>

    > Only admin user allowed 
            """)
        with gr.Row():
            password = gr.Textbox(label='Enter the Password')
            collection_name = gr.Textbox(label='Collection Name')
            page_start = gr.Textbox(label='Page Start')
        file_type = gr.Dropdown(['PDF', 'TEXT', 'STRING'], label='Select File Type',
                               value = 'PDF')
        file_fields = gr.File(visible = True, interactive=True)
        context = gr.Textbox(label="Enter the Context", visible = False)
        btn = gr.Button("Submit")

        btn.click(embed_and_store, inputs=[password, collection_name, file_type, file_fields, context,page_start], outputs=[file_fields, context])
        file_type.change(update_interface, inputs=[file_type], outputs=[context, file_fields])
    gr.Markdown("""
    <div align='center'>It could be helpful for making RAG application</div>
               <div align='center'>| MONGODB | LANGCHAIN | HUGGINGFACE | MITSRAL |</div>
            """)    

demo.launch()