import gradio as gr
from rag import mongo_rag_tool
from gradio.themes.base import Base

# Create an instance of GradIO
with gr.Blocks(theme=Base(), title="Market Research and VOC bot") as demo:
    # A styled header for the app
    gr.Markdown(
        """
        <div style='text-align: center; font-size: 24px; font-weight: bold; margin-bottom: 20px;'>
        Chat with your data
        </div>
        """
    )
    
    # Input fields for the collection and the question, with descriptive text
    gr.Markdown(
        """
        <div style='text-align: left; font-size: 18px; margin-bottom: 10px;'>
        Enter the collection and your query to get relevant answers:
        </div>
        """
    )
    collection_textbox = gr.Textbox(label="Enter your Collection:", placeholder="e.g., market_data", lines=1)
    query_textbox = gr.Textbox(label="Enter your Question:", placeholder="Type your question here...", lines=1)

    # Submit button with some spacing and central alignment
    with gr.Row():
        button = gr.Button("Submit", variant="primary", size="lg")

    # Output section for displaying the answer and sources one below the other
    gr.Markdown(
        """
        <div style='text-align: left; font-size: 18px; margin-bottom: 10px;'>
        Output:
        </div>
        """
    )
    # Using a Column to place the answer and source outputs one after the other
    with gr.Column():
        answer_output = gr.Textbox(lines=5, label="Answer:", max_lines=50)
        source_output = gr.Textbox(lines=5, label="Sources:", max_lines=50)

    # Connect the button to the function
    button.click(mongo_rag_tool, inputs=[query_textbox, collection_textbox], outputs=[answer_output, source_output])

demo.launch()