import gradio as gr

from utils import rag_api, agent_api, base64_to_audio


# Helper functions
def assignment_part_1(input_text, language):
    response = rag_api(input_text, language)
    if response["audio"]:
        output_audio = base64_to_audio(response["audio"])
        return response, output_audio
    else:
        return response, None, None


def assignment_part_2(input_text, language):
    response = agent_api(input_text, language)
    if response["audio"]:
        output_audio = base64_to_audio(response["audio"])
        return response, output_audio
    else:
        return response, None, None


# Create the Gradio interface
with gr.Blocks() as block:

    gr.Markdown("# <p style='text-align:center;'>Sarvam Intern Assignment</p>")

    # User Query
    input_text = gr.Textbox(label="Enter your query:")

    # Language selection
    language = gr.Dropdown(label="Select audio_language_code",
                           choices=["hi-IN", "bn-IN", "kn-IN", "ml-IN",
                                    "mr-IN", "od-IN", "pa-IN", "ta-IN",
                                    "te-IN", "gu-IN", "en-IN"],)

    # Two buttons
    with gr.Row():
        button_1 = gr.Button("RAG")
        button_2 = gr.Button("Agent + RAG")

    # Output field
    output_text = gr.JSON(label="Output")
    output_audio = gr.Audio(type="numpy",
                            label="Output Audio",
                            interactive=False)

    # Button click actions
    button_1.click(fn=assignment_part_1,
                   inputs=[input_text, language],
                   outputs=[output_text, output_audio])

    button_2.click(fn=assignment_part_2,
                   inputs=[input_text, language],
                   outputs=[output_text, output_audio])

# Launch the Gradio app
block.launch()