import gradio as gr
from openai import OpenAI
import os


client = OpenAI(api_key=os.environ["API_KEY"])

def play_tts(x: gr.LikeData):
    response = client.audio.speech.create(
        model="tts-1",
        voice="alloy",
        input=x.value
    )
    response.stream_to_file('temp.mp3')
    f = gr.Audio('temp.mp3', streaming=True, elem_id='audio')
    print(f'"{x.value}" saved temp.mp3 file')

    return f


def answer(state, state_chatbot, text):
    messages = state + [{
        'role': 'user',
        'content': text
    }]

    res = client.chat.completions.create(
        model='gpt-4',
        messages=messages
    )

    msg = res.choices[0].message.content

    new_state = [{
        'role': 'user',
        'content': text
    }, {
        'role': 'assistant',
        'content': msg
    }]

    state = state + new_state
    state_chatbot = state_chatbot + [(text, msg)]

    print(state)

    return state, state_chatbot, state_chatbot


with gr.Blocks() as demo:
    with gr.Row():
        gr.HTML("""<div style="text-align: center; max-width: 500px; margin: 0 auto;">
            <div>
                <h1>Dongwook's ChatBot with TTS</h1>
                <h2>재생을 원하는 부분을 클릭하세요 </h2>
            </div>
        </div>""")
    with gr.Row():
        audios = gr.Audio('temp.mp3', streaming=True, elem_id='audio')

    state = gr.State([{
        'role': 'system',
        'content': 'You are a helpful assistant.'
    }])
    state_chatbot = gr.State([])

    with gr.Row():
        chatbot = gr.Chatbot(
            [],
            elem_id="chatbot",
            bubble_full_width=False,
            show_copy_button=True,
            likeable=False,
        )
        chatbot.select(play_tts, None, outputs=audios)
    # chatbot.click(fn=print_play, inputs=None,outputs=None)

    with gr.Row():
        txt = gr.Textbox(
            scale=4,
            show_label=False,
            placeholder="Enter text and press enter",
            container=False,
        )

    txt_msg = txt.submit(answer, [state, state_chatbot, txt], [state, state_chatbot, chatbot])
    txt_msg.then(lambda: '', None, txt, queue=False)

demo.queue()
if __name__ == "__main__":
    demo.launch(share=True)