import api
import gradio as gr
import random

with gr.Blocks() as demo:
    with gr.Column() as setup_col:
        gr.HTML("<h1 style='text-align: center; margin-top: 1em; font-size: 2em;'>Rizz++</h1>")
        gr.Markdown("Rizz++ simulates conversations you might have when approaching someone you're interested in, helping you develop your skills. \nRandomly generate a scenario involving...")
        with gr.Row():
            approach_girl = gr.Button("Approaching a Girl")
            approach_boy = gr.Button("Approaching a Boy")
        gr.Markdown("...or enter your own scenario:")
        with gr.Group():
            custom_scenario = gr.Textbox(lines=3, label="Scenario", placeholder="There's a cute girl in my boxing class I want to talk to.")
            scenario_btn = gr.Button("Submit Scenario")

    with gr.Column(visible=False) as convo_col:
        chat = gr.Chatbot()
        recording = gr.Audio(label="Your Line", sources="microphone", type="filepath")
        with gr.Group():
            with gr.Row():
                help_btn = gr.Button("Help Me", scale=0)
                suggestion = gr.Textbox(container=False, placeholder="Stuck? Let Rizz++ suggest your next line!")

    @gr.on([approach_boy.click, approach_girl.click, scenario_btn.click], outputs=[setup_col, convo_col])
    def start_convo():
        return {
            setup_col: gr.Column(visible=False),
            convo_col: gr.Column(visible=True)
        } 
    
    scenario = gr.State()
    personality = gr.State(lambda: random.choice(["shy", "confident", "awkward", "charming", "rude"]))
    interested = gr.State(lambda: random.choice([True, False]))
    SCENARIO_LINE = "The scenario is: **{}** \n\n Record your opening line below!"

    def start_chat(scenario):
        return [[SCENARIO_LINE.format(scenario), None]]

    approach_boy.click(lambda: api.generate_scenario("boy"), outputs=scenario).then(start_chat, scenario, chat)
    approach_girl.click(lambda: api.generate_scenario("girl"), outputs=scenario).then(start_chat, scenario, chat)
    scenario_btn.click(lambda x: x, custom_scenario, scenario).then(start_chat, scenario, chat)


    def transcribe(recording, chat):
        user_message = api.transcribe_audio(recording)
        chat.append([user_message, None])
        return chat
    
    def respond(scenario, chat, personality, interested):
        response = api.generate_response(scenario, chat[1:], personality, interested)
        chat[-1][1] = response
        return chat


    recording.stop_recording(transcribe, inputs=[recording, chat], outputs=[chat]
    ).then(respond, inputs=[scenario, chat, personality, interested], outputs=[chat]
    ).then(lambda: None, outputs=[recording])

    help_btn.click(lambda scenario, chat: api.suggest_next_line(scenario, chat[1:]), inputs=[scenario, chat], outputs=[suggestion])
    
demo.launch()