import os
import gradio as gr
from gradio.components import Textbox, Button, Slider, Checkbox
from AinaTheme import theme

from vectorstore import VectorStore

MAX_NEW_TOKENS = 700


vectorStore = VectorStore(embeddings_model=os.getenv("EMBEDDINGS", "BAAI/bge-m3"))

def eadop_rag(prompt, num_chunks):
    prompt = prompt.strip()
    if prompt == "":
        gr.Warning(
            "Prompt can't be empty!"
        )
        raise ValueError("Prompt can't be empty!")
    return vectorStore.get_context(prompt, num_chunks)
    
def clear():
    return (
        None,
        None,
        gr.Slider(value=2.0),
    )

def gradio_app():
    with gr.Blocks(theme=theme) as demo:
        with gr.Row(equal_height=True):
            output = Textbox(
                lines=10, 
                label="Context", 
                interactive=False, 
                show_copy_button=True
            )

        with gr.Row(equal_height=True):
            input_ = Textbox(
                label="Input",
                placeholder="Quina és la finalitat del Servei Meteorològic de Catalunya?",
            )

        with gr.Row(equal_height=True):
            num_chunks = Slider(
                            minimum=1,
                            maximum=6,
                            step=1,
                            value=2,
                            label="Number of chunks"
                        )
        with gr.Row(equal_height=True):
            clear_btn = Button("Clear")
        with gr.Row(equal_height=True):
            submit_btn = Button("Submit", variant="primary")
                        

        clear_btn.click(
            fn=clear, 
            inputs=[], 
            outputs=[input_, output, num_chunks],
              queue=False, 
              api_name=False
        )
        
        submit_btn.click(
            fn=eadop_rag, 
            inputs=[input_, num_chunks], 
            outputs=[output],
            api_name="get-eadop-rag"
        )

        demo.launch(show_api=True)


if __name__ == "__main__":
    gradio_app()