import gradio as gr
import time
from utils import *
import os

grammar_prompt="""
I want you to act as a grammar mistake checker and make the sentence more fluent. You take all the user input and auto correct it. Just reply to user input with correct grammar and reasons, DO NOT reply the context of the question of the user input. If the user input is grammatically correct and fluent, just ignore it. For each wrong sentence, you will show Original, Corrected, and Reason. Sample of the conversation will show below: 

Correct: today is a good day.

Original: today is a good day.
Corrected: Today is a good day. 
Reason: Capitalize the first letter of the sentence.

###

"""
def main():        
    chat_history = [
        {"role": "system", "content": os.environ.get("SECRET_PROMPT","You are a chat bot. Talk to me!")},
    ]

    def init():
        nonlocal chat_history
        chat_history = [
            {"role": "system", "content": os.environ.get("SECRET_PROMPT","You are a chat bot. Talk to me!")},
        ]
    init()
    
    def convert_chatbox(chat_history):
        return [f"{i['role']}: {i['content']}" for i in chat_history]

    with gr.Blocks() as demo:
        chatbot = gr.Chatbot()
        with gr.Row():
            msg = gr.Textbox()
            audio = gr.Audio(source="microphone", type="filepath", streaming=False)
        player = gr.Audio( type="filepath", label="Speaker",interactive=False)
        summary = gr.Button("Summary")
        summary_box = gr.Textbox(label="Summary")
        # functions
        def respond(message):
            # TODO: replace this with real GPT model
            chat_history.append({'role': 'user', 'content': message})
            result = generate_response(chat_history)
            mesg=result['choices'][0]['message']
            print("recv: ", mesg)
            
            response = mesg['content']
            chat_history.append(mesg)
            
            # write to file
            result = tts(response)
            with open("/tmp/temp.wav", "wb") as audio_file:
                audio_file.write(result.audio_data)

            print("write to temp.wav")
            
            chatbot.value.append((message,response))
            print("chat_history: ", chatbot.value)
            
            return None, "/tmp/temp.wav", chatbot.value

        msg.submit(respond, [msg], [msg, player,chatbot])
        
        
        def transcribe(audio_file):
            print("start transcribe, ", audio_file)
            
            start = time.time()
            text = recognize_from_file(audio_file)
            print("use ", time.time()-start)
            
            print("transcribe done, ", text)
            return respond(text)

        audio.change(transcribe, [audio], [audio, player, chatbot])
        
        def summary_response():
            messages = [
            ]
            sentences = []
            for user,assistant in chatbot.value:
                sentences.append("Correct: " + user)
            messages.append({'role': 'user', 'content': grammar_prompt + "\n".join(sentences)})
            
            result = generate_response(messages)
            mesg=result['choices'][0]['message']
            corrected = mesg['content']
            print("recv: ", mesg)
            
            return corrected
        
        summary.click(summary_response, None, summary_box, queue=False)


        btn = gr.Button("导出 & 重置", type="button", label="导出 & 重置")
        outputs = gr.JSON()

        def export():
            stats_history = {}
            for i,item in enumerate(chatbot.value):
                user,assistant = item
                stats_history[str(i)] = {
                    "user": user,
                    "assistant": assistant
                }
            init()
            chatbot.value = []
            return  [], stats_history

        btn.click(export, None, [ chatbot, outputs])

    demo.launch(show_error=True)
    
if __name__ == "__main__":
    main()