import gradio as gr

from main import (
    load_dataset,
    build_vector_db,
    generate_response
)

DATA_FILE = "cat-facts.txt"


def initialize_rag():
    dataset = load_dataset(DATA_FILE)
    build_vector_db(dataset)


# Initialize RAG once when the app starts
initialize_rag()


def answer_question(query):
    if not query.strip():
        return "Please enter a question."

    answer, results = generate_response(query)
    return answer


with gr.Blocks(title="Local RAG Chatbot for Cat Facts") as demo:
    gr.Markdown("# Local RAG Chatbot for Cat Facts")

    query = gr.Textbox(
        label="Ask a question about cats",
        placeholder="e.g. What is the average lifespan of a cat?"
    )

    answer = gr.Textbox(
        label="Answer",
        lines=5
    )

    submit_btn = gr.Button("Answer")

    submit_btn.click(
        fn=answer_question,
        inputs=query,
        outputs=answer
    )

demo.launch()