from pathlib import Path

import gradio as gr
import librosa
import time
from loguru import logger
from transformers import pipeline

logger.info("Loading the ASR pipeline...")
pipe = pipeline(
    "automatic-speech-recognition",
    model="Parakeet-Inc/furigana_whisper_small_jsut",
)
logger.success(f"ASR pipeline loaded successfully, device: {pipe.device}")


def transcribe_with_prompt(audio_path: str | Path, prompt: str) -> str:
    audio_path = Path(audio_path)
    logger.info(f"Got audio: {audio_path.name}, prompt: {prompt}")
    duration = librosa.get_duration(path=audio_path)
    logger.info(f"Audio duration: {duration:.2f} seconds")
    if duration > 30:
        raise gr.Error(
            "音声ファイルの長さが30秒を超えています。30秒以内の音声ファイルを使用してください。"
        )
    t0 = time.perf_counter()
    prompt_ids = pipe.tokenizer.get_prompt_ids(prompt, return_tensors="pt").to(
        pipe.device
    )
    generate_kwargs = {"prompt_ids": prompt_ids}
    result = pipe(str(audio_path), generate_kwargs=generate_kwargs)["text"]
    t1 = time.perf_counter()
    logger.success(f"Transcription completed in {t1 - t0:.2f} sec:\n{result}")
    return result


with gr.Blocks() as app:
    gr.Markdown("# ふりがなWhisperデモ")
    gr.Markdown(
        "音声と、音声の漢字仮名交じり文を入力すると、発話されている読みのカタカナ列を出力します。Smallモデルなのでそこまで精度がよくないかもしれません。"
    )
    gr.Markdown("また、デモではCPU使用のため極めて遅いです。モデルのREADMEや記事を参考にしてローカルで動かすことを推奨します。")
    gr.Markdown("モデル: [Parakeet-Inc/furigana_whisper_small_jsut](https://huggingface.co/Parakeet-Inc/furigana_whisper_small_jsut)")
    gr.Markdown("ライセンス: Apache-2.0")

    with gr.Row():
        audio_input = gr.Audio(label="音声ファイル", type="filepath")
        prompt_input = gr.Textbox(
            label="漢字仮名交じり文（プロンプト）",
            placeholder="音声で発話されている内容を漢字仮名交じり文で入力",
        )

    transcribe_button = gr.Button("Transcribe")
    transcription_output = gr.Textbox(label="結果")

    transcribe_button.click(
        fn=transcribe_with_prompt,
        inputs=[audio_input, prompt_input],
        outputs=transcription_output,
    )

app.launch()