import gradio as gr
from src.rag import RAG

def upload_files(files, filepaths):
    verbose = False
    filepaths_new = [file.name for file in files]
    if verbose:
        print(f'previous files: {filepaths}')
        print(f'new files: {filepaths_new}')
    filepaths = filepaths + filepaths_new
    return filepaths, filepaths

def initialize_rag(pdfs):
    print(f'Initializing RAG-Chatbot with:\npdfs:\n{pdfs}')
    # Initialize RAG instance
    try:
        rag = RAG(
            urls=[],
            pdfs=pdfs,
            k=2
            )
        message = "RAG-Chatbot initialized successfully!"
    except Exception as e:
        rag = None
        message = f"Error initializing RAG-Chatbot:\n{e}"
    return message, rag

def get_rag_response(message, history, rag):
    if rag is None:
        return "Error: RAG-Chatbot is not initialized yet!"
    print(f"Question: {message}")
    response = rag.ask_QAbot(message)
    answer_str = response['answer']
    sources = [f"{i+1}. {source.split('/')[-1]}" for i, source in enumerate(response['sources'])]
    sources_str = ';'.join(sources)
    print(sources_str)
    response_str = f"""
    {answer_str}
    
    Sources: 
    {sources_str}
    """

    return response_str

with gr.Blocks() as demo:
    gr.Markdown("# RAG-Chatbot")
    gr.Markdown("## Instructions")
    gr.Markdown("""
                Upload PDF's that will form the basis of the database for our RAG-Chatbot. 
                Once done adding documents (multiple uploads allowed), click `Initialize` to start the building process.
                Note that building time depends on the number and length of uploaded documents.
                When the building is done, 
                as will be indicated by `RAG-Chatbot initialized successfully!` in the Initialization Status box, 
                you can start chatting!
                """)
    # PDFs
    gr.Markdown("## 1. Build the bot")
    pdfpaths = gr.State([])
    file_output = gr.File()
    upload_button = gr.UploadButton("Upload PDF(s)",  file_count="multiple")
    upload_button.upload(
        fn=upload_files, 
        inputs=[upload_button, pdfpaths], 
        outputs=[file_output, pdfpaths])
    # State to store the RAG instance
    rag_instance = gr.State(None)  # Initially None
    init_button = gr.Button("Initialize")
    init_status = gr.Textbox(label="Initialization Status", interactive=False)
    # Event handlers
    init_button.click(
        initialize_rag,
        inputs=[pdfpaths],
        outputs=[init_status, rag_instance]  # Output: status message and the RAG instance
    )
    gr.Markdown('## 2. Chat')
    # Chat Interface for RAG-Chatbot
    gr.ChatInterface(
        fn=get_rag_response,
        additional_inputs=[rag_instance],
        type="messages"
    )

if __name__ == "__main__":
    demo.launch()