import gradio as gr
from openai import OpenAI
import os
import time

audio_format = [".mp3", ".mp4", ".mpeg", ".mpga", ".m4a", ".wav", ".webm"
                ,".MP3", ".MP4", ".MPEG", ".MPGA", ".M4A", ".WAV", ".WEBM"]

lang_code = {'Japanese': "ja", 'English': "en"}
mode_code = {'文字': "transcribe", '翻訳': "translate"}

max_file_size = 25 * 1024 * 1024

def set_state(openai_key, lang, mode, state):

    state["openai_key"]= openai_key
    state["lang"] = lang_code[lang]
    state["mode"] = mode_code[mode[0:2]]

    return state

def create_textfile(voice_msg, up_file ,state):

    # OpenAIキーチェック
    if state["openai_key"] == "":

        err_msg = "OpenAIキーを入力してください。（設定タブ）"

        return voice_msg, up_file, "", None, err_msg

    # ファイル入力チェック
    if up_file is None and voice_msg is None:

        err_msg = "マイクかファイルで音声を入力してください。"

        return None, None, "", None, err_msg

    # アップロードの場合ファイルチェック
    if voice_msg is None:

        # ファイルパスなど取得
        file_name = os.path.splitext(os.path.basename(up_file.name))[0]
        root, ext = os.path.splitext(up_file.name)

        if ext not in audio_format:

            # ファイル形式チェック
            err_msg = "指定した形式のファイルをアップしてください。（mp3, mp4, mpeg, mpga, m4a, wav, webm）"

            return None, None, "", None, err_msg

        file_size = os.path.getsize(up_file.name)

        if file_size >= max_file_size:

            # ファイルサイズチェック
            err_msg = "ファイルが大きすぎます。25MB未満にして下さい。"

            return None, up_file.name, "", None, err_msg

    # 音声ファイルセット
    if voice_msg:
        audio_file= open(voice_msg, "rb")

    else:
        audio_file= open(up_file.name, "rb")

    try:

      # whisperで変換処理
      os.environ["OPENAI_API_KEY"] = state["openai_key"]

      # クライアント新規作成
      client = OpenAI()

      # client作成後は消す
      os.environ["OPENAI_API_KEY"] = ""

      if state["mode"] == "transcribe":

          # 文字起こし実行
          trans_text = client.audio.transcriptions.create(model="whisper-1", file=audio_file ,language = state["lang"], response_format="text")

      else:

          # 翻訳文字起こし実行
          trans_text = client.audio.translations.create(model="whisper-1", file=audio_file, response_format="text")

      # ファイル名設定
      if voice_msg:
          trans_file = trans_text[:8] + "_whisper.txt"
      else:
          trans_file = file_name + "_whisper.txt"

      with open(trans_file, mode="w") as f:

          # テキストに書き出す
          f.write(trans_text)

      return None, None, trans_text, trans_file, ""

    except Exception as e:

      return voice_msg, None, "", None, e

def wait_mic_process():

    # 録音後処理に時間がかかるので待つ
    time.sleep(1)

    return gr.update(interactive=True)

with gr.Blocks() as demo:

    title = "<h2>Whisperデモアプリ【基本版】</h2>"
    message = "<h3>最初に[設定]タブからOpenAIキーを入力してください。"
    message += "</h3>"

    gr.Markdown(title + message)

    # セッションの宣言
    state = gr.State({
        "openai_key" : "",
        "lang": "",
        "mode" : "",
    })

    with gr.Tab("whisperを利用する") as maintab:

      # 各コンポーネント定義
      voice_msg=gr.components.Audio(sources="microphone",type="filepath", label="音声入力")
      up_file = gr.File(file_types=[".mp3", ".mp4", ".mpeg", ".mpga", ".m4a", ".wav", ".webm"], label="音声ファイルアップロード", type="filepath")

      # ボタン類
      with gr.Row():
        btn = gr.Button("文字に起こす")
        clear = gr.ClearButton([voice_msg, up_file], value="クリア")

      # 出力
      sys_msg = gr.Text(label="システムメッセージ")
      text = gr.TextArea(label="文字起こし内容")
      file = gr.File(label="出力テキストファイル")

      # 送信ボタンクリック時の処理
      btn.click(create_textfile, inputs=[voice_msg, up_file, state], outputs=[voice_msg, up_file, text, file, sys_msg], queue=True)

      # 音声録音後処理に時間がかかるのでボタンを利用不可にする
      voice_msg.start_recording(lambda:gr.update(interactive=False), None, btn)
      voice_msg.stop_recording(wait_mic_process, None, btn)

    with gr.Tab("設定") as settab:
      openai_key = gr.Textbox(label="OpenAI API Key", interactive = True)
      lang = gr.Dropdown(label="Language", choices=["Japanese", "English"], value = "Japanese", interactive = True)
      mode = gr.Dropdown(label="Mode", choices=["文字起こし", "翻訳＋文字起こし"], value = "文字起こし", interactive = True)

      # 設定変更時
      maintab.select(set_state, [openai_key, lang, mode, state], state)

    with gr.Tab("利用上の注意"):

      gr.Markdown("・マイクの音声とファイル入力の両方がある場合はマイクが優先されます。")
      gr.Markdown("・文字に起こせるファイル形式は mp3, mp4, mpeg, mpga, m4a, wav, webm のみです。")
      gr.Markdown("・翻訳は日本語の音声を英語のテキストに変えます。英語の音声→日本語はできません。")
      gr.Markdown("※翻訳機能はマイク入力の精度が低いです。ファイルでの入力をおすすめします。")

demo.queue()
demo.launch(debug=False)