import os
import sys

from shiny.express import input, ui

from all_rag_fns import do_rag

oai_api_key = os.getenv("OPENAI_API_KEY")
ui.page_opts(
    title="Use Shiny to Run RAG on the previous R/Gov Talks",
    fillable=True,
    fillable_mobile=True,
)


with ui.layout_sidebar():
    # Add radio buttons in the sidebar
    with ui.sidebar():
        ui.input_radio_buttons(
            "model_choice",
            "Select Model:",
            choices={"gpt-4o-mini": "Cheaper", "gpt-4o": "More Accurate"},
            selected="gpt-4o-mini",
        )

    # Create a chat instance and display it in the main panel
    chat = ui.Chat(id="chat")
    chat.ui()


# Define a callback to run when the user submits a message
@chat.on_user_submit
async def _():
    user_message = chat.user_input()
    response, _ = do_rag(
        user_input=user_message,
        n_results=3,
        stream=True,
        oai_api_key=oai_api_key,
        model_name=input.model_choice(),
    )
    # Append the response into the chat
    await chat.append_message_stream(response)