import os
import gradio as gr
from openai import OpenAI
from yt_dlp import YoutubeDL
from pydub import AudioSegment
import math

# 定数
MAX_LENGTH = 5000
MODEL_SUMMARY = "gpt-3.5-turbo-0125"

# 分割時間設定（20分）
split_time = 20 * 60 * 1000

lang_code = {'Japanese': "ja", 'English': "en"}

def set_state(openai_key, lang, state):
    """ 設定タブの情報をセッションに保存する関数 """

    state["openai_key"]= openai_key
    state["lang"] = lang_code[lang]

    return state

def youtube_mp3(url):

  # ファイル名は動画のIDに
  mp3_name = url[url.rfind('/') + 1:].replace("watch?v=","")


  # フォーマット設定
  ydl_opts = {
      'outtmpl': mp3_name,
      "format": "mp3/bestaudio/best",
      "postprocessors": [
          {
              "key": "FFmpegExtractAudio",
              "preferredcodec": "mp3",
          }
      ],
  }

  try :
      # yt-dlpでmp3に
      with YoutubeDL(ydl_opts) as ydl:
          result = ydl.download([url])

      return mp3_name + ".mp3"

  except Exception as e:
      print(e)
      return None


def create_textfile(url, file, state):

    err_msg = ""

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

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

        return None, err_msg

    # URL入力チェック
    if url == "" and file is None:

        err_msg = "URLを入力するか、音声ファイルをアップして下さい。"

        return None, err_msg

    # OpenAIクライアント作成
    os.environ["OPENAI_API_KEY"] = state["openai_key"]

    client = OpenAI()

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

    # 動画を音声ファイルにする
    if url != "":

        file_name = youtube_mp3(url)

        if file_name is None:

            err_msg = "音声ファイルにする作業でエラーが発生しました。URLを確認して下さい。"

            return None, err_msg

    else:
        file_name = file

    # 音声ファイルを分割
    audio_list = audio_seg(file_name)

    # whisperで文字に起こす
    text_list = whisper_audio(client, audio_list, state["lang"])

    if text_list is None:

        err_msg = "whisperでエラーが発生しました。OpenAI APIキーが正しいか、クレジット残高があるか確認して下さい。"

        return None, err_msg


    # テキストファイルを結合
    text_name = file_name.replace(".mp3",".txt")

    combin_text(text_list, text_name)

    return text_name, err_msg


def whisper_audio(client, audio_list, lang):

    text_list = []

    try:

        for audio in audio_list:

              audio_file= open(audio, "rb")

              trans_text = client.audio.transcriptions.create(model="whisper-1", file=audio_file, language=lang, response_format="text")

              # ファイル名設定
              text_name = audio.replace(".mp3","_.txt")

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

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

              text_list.append(text_name)

        return text_list

    except Exception as e:

        return None


def audio_seg(file_name):

    # 分割したリスト
    div_file = []

    # ファイルから音声情報取得
    audio = AudioSegment.from_mp3(file_name)

    # 音声の長さ取得（ミリ秒）
    dur_mlseconds = audio.duration_seconds * 1000

    # 分割数を決める
    div_count = math.ceil(dur_mlseconds / split_time)

    if div_count == 1:

        # 分割なしの時は元ファイルのみ
        div_file.append(file_name)

        return div_file

    # 分割時間初期設定
    start = 0
    end = split_time

    for i in range(div_count):

        div_audio = audio[start:end]
        div_audio.export(str(i+1) + "_" + file_name, format="mp3")

        # ファイル名をセット
        div_file.append(str(i+1) + "_" + file_name)

        start = end
        end += split_time

    return div_file


def combin_text(text_list, text_name):

    # ファイルを一つにまとめる
    with open(text_name, "w") as of:
        for file in text_list:
            with open(file, "r") as f:
                of.write(f.read())

                # 必要であれば改行を加える
                # outfile.write(infile.read() + "\n")

    return

def create_mp3(url, state):

    err_msg = ""

    # URL入力チェック
    if url == "":

        err_msg = "URLを入力して下さい。"

        return None, err_msg

    # 動画を音声ファイルにする
    file_name = youtube_mp3(url)

    if file_name is None:

        err_msg = "音声ファイルにする作業でエラーが発生しました。URLを確認して下さい。"

        return None, err_msg

    return file_name, ""

def create_summary(file, state):

    print(file)

    err_msg = ""

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

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

        return None, err_msg

    # URL入力チェック
    if file is None:

        err_msg = "テキストがありません。"

        return None, err_msg

    # OpenAIクライアント作成
    os.environ["OPENAI_API_KEY"] = state["openai_key"]

    client = OpenAI()

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

    summary, err_msg = exec_summary(client, file)

    if err_msg != "":

        return file, err_msg

    sum_file = "summary_" + os.path.basename(file)

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

        f.write(summary)

    return sum_file, ""


def exec_summary(client, file):

    try:

        with open(file, 'r') as f:

            text = f.read()

        if len(text) > MAX_LENGTH:

            err_msg = "要約の文字数上限を超えています。"

            return "", err_msg

        messages=[
              {"role": "system", "content": "あなたは優秀なアシスタントです。与えられた文章を要約して下さい。"},
              {"role": "user", "content": text},
          ]

        # GPTに問い合わせ
        response = client.chat.completions.create(
            model=MODEL_SUMMARY,
            messages=messages,
            # max_tokens=MAX_TOKENS,
        )

        summary = response.choices[0].message.content

        return summary, ""

    except Exception as e:
        print(e)
        return "", "要約作成でエラーが発生しました。"


with gr.Blocks() as demo:

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

    gr.Markdown(title + message)

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

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

      # 各コンポーネント定義
      url = gr.Text(label="YouTubeのURL")

      with gr.Accordion(label="音声ファイルをアップする", open=False):
          up_file = gr.File(file_types=[".mp3", ".mp4", ".mpeg", ".mpga", ".m4a", ".wav", ".webm"], label="音声ファイルアップロード")

      # ボタン類
      with gr.Row():
          btn_txt = gr.Button("テキスト作成")
          btn_mp3 = gr.Button("MP3作成")
          # clear = gr.ClearButton(value="リセット", components=[url, sys_msg])

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

      with gr.Accordion(label="テキストを要約する", open=False):
          btn_sum = gr.Button("テキスト要約")
          sum_file = gr.File(label="要約テキストファイル", interactive = False)

      # 送信ボタンクリック時の処理
      btn_txt.click(create_textfile, inputs=[url, up_file, state], outputs=[out_file, sys_msg], queue=False)
      btn_mp3.click(create_mp3, inputs=[url, state], outputs=[out_file, sys_msg], queue=False)
      btn_sum.click(create_summary, inputs=[out_file, state], outputs=[sum_file, sys_msg], queue=False)

    with gr.Tab("設定") as set:
      openai_key = gr.Textbox(label="OpenAI API Key", interactive = True)
      lang = gr.Dropdown(choices=["Japanese", "English"], value = "Japanese", label="Language", interactive = True)

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

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

      caution = '・URLとファイルが両方ある場合はURLが優先されます。<br>'
      caution += "・文字起こしにはお金がかかりますが（1分あたり約0.9円）、MP3作成のみは無料です。<br>"
      caution += "・要約の際のGPTのモデルはGPT3.5です（gpt-3.5-turbo-0125）<br>"
      caution += "・要約には料金が発生します。（1000文字あたり約0.2円）<br>"
      gr.Markdown("<h3>" + caution + "</h3>")

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