import gradio as gr
from pipeline import preprocessing_pipeline, conversational_rag
from pipeline import system_message, user_message
from haystack.dataclasses import ChatMessage
import time
import os

def process_files_into_docs(files, progress=gr.Progress()):
    if isinstance(files, dict):
        files = [files]
    
    if not files:
        return 'No file uploaded!'
        
    preprocessing_pipeline.run({'file_type_router': {'sources': files}})
    return "Database created🤗🤗"

def rag(history, question):
    if history is None:
        history = []

    # Run Haystack pipeline
    res = conversational_rag.run(
        data={
            "query_rephrase_prompt_builder": {"query": question},
            "prompt_builder": {
                "template": [system_message, user_message],
                "query": question,
            },
            "memory_joiner": {
                "values": [ChatMessage.from_user(question)]
            }
        },
        include_outputs_from=["llm"]
    )

    bot_message = res["llm"]["replies"][0].content

    # Add user message
    history = history + [
        {"role": "user", "content": question}
    ]

    # Stream assistant message
    streamed = ""
    for token in bot_message.split():
        streamed += token + " "
        yield (
            history + [{"role": "assistant", "content": streamed.strip()}],
            ""
        )
        time.sleep(0.05)

    # Final assistant message
    history = history + [
        {"role": "assistant", "content": bot_message}
    ]

    yield history, ""


EXAMPLE_FILE = "RAG Survey.pdf"

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("<center><h1>TalkToFiles - Query your documents! 📂📄</h1></center>")
    gr.Markdown("""##### This AI chatbot🤖 can help you chat with your documents. Can upload <b>Text(.txt), PDF(.pdf) and Markdown(.md)</b> files.
    <b>Please do not upload confidential documents.</b>""")
    
    with gr.Row():
        with gr.Column(scale=86):
            gr.Markdown("""#### ***Step 1 - Upload Documents and Initialize RAG pipeline***</br> Can upload Multiple documents""")
            
            with gr.Row():
                file_input = gr.File(
                    label='Upload Files',
                    file_count='multiple',
                    file_types=['.pdf', '.txt', '.md'],
                    interactive=True
                )
                
            with gr.Row():
                process_files = gr.Button('Create Document store')
                
            with gr.Row():
                result = gr.Textbox(label="Document store", value='Document store not initialized')
                
            # Pre-processing Events
            process_files.click(
                fn=process_files_into_docs,
                inputs=file_input,
                outputs=result,
                show_progress=True
            )
            
            # def load_example():
            #     return [EXAMPLE_FILE]
                
            # with gr.Row():
            #     gr.Examples(
            #         examples=[[EXAMPLE_FILE]],
            #         inputs=file_input,
            #         examples_per_page=1,
            #         label="Click to upload an example"
            #     ).dataset.click(fn=load_example, inputs=[], outputs=file_input)
                
        with gr.Column(scale=200):
            gr.Markdown("""#### ***Step 2 - Chat with your docs*** """)
            chatbot = gr.Chatbot(label='ChatBot', type="messages")  # <-- Added type="messages" to fix deprecation
            user_input = gr.Textbox(label='Enter your query', placeholder='Type here...')
            
            with gr.Row():
                submit_button = gr.Button("Submit")
                clear_btn = gr.ClearButton([user_input, chatbot], value='Clear')
                
            submit_button.click(
                rag,
                inputs=[chatbot, user_input],
                outputs=[chatbot, user_input]
            )

# Use api_name=None to avoid API generation issues
demo.launch()