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 = []
        
    messages = [system_message, user_message]
    res = conversational_rag.run(
        data = {
            'query_rephrase_prompt_builder' : {'query': question},
            'prompt_builder': {'template': messages, 'query': question},
            'memory_joiner': {'values': [ChatMessage.from_user(question)]}
        },
        include_outputs_from=['llm', 'query_rephrase_llm']
    )
    
    bot_message = res['llm']['replies'][0].content
    streamed_message = ""
    
    for token in bot_message.split():
        streamed_message += f"{token} "
        yield history + [(question, streamed_message.strip())], " "
        time.sleep(0.05)
        
    history.append((question, 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()