import gradio as gr
import whisper
import tempfile
import librosa
import soundfile as sf
import os

# 檢查是否有可用的 GPU，否則使用 CPU
device = "cuda" if whisper.torch.cuda.is_available() else "cpu"

# 載入 Whisper 模型並指定設備（GPU 或 CPU）
model = whisper.load_model("small", device=device)

def transcribe_audio_and_save(audio):
    # 根據上傳的音效檔路徑獲取檔名
    input_filename = os.path.splitext(os.path.basename(audio))[0]
    
    # 將音效檔轉換為 WAV 格式並處理
    with tempfile.NamedTemporaryFile(suffix=".wav") as audio_file:
        # 使用 librosa 讀取音效檔，支援多種音效格式
        audio_data, sample_rate = librosa.load(audio, sr=16000)
        sf.write(audio_file.name, audio_data, sample_rate)

        # 使用 Whisper 辨識音訊並啟用時間戳功能
        result = model.transcribe(audio_file.name, task='transcribe', verbose=True)
        original_text = result['text']  # 直接的 Whisper 輸出
        segments = result['segments']  # 使用段落資訊
        
        # 將辨識結果每段之間加入空格
        text_with_spaces = ""
        for segment in segments:
            text_with_spaces += segment['text'] + " "  # 在每句之間加入空格
        
        # 儲存辨識結果到文字檔，檔名根據音效檔檔名生成
        output_filename = f"{input_filename}_transcription.txt"
        output_path = os.path.join(tempfile.gettempdir(), output_filename)
        with open(output_path, "w") as f:
            f.write("Original transcription:\n")
            f.write(original_text + "\n\n")
            f.write("Transcription with spaces:\n")
            f.write(text_with_spaces.strip() + "\n")
        
    return original_text, text_with_spaces.strip(), output_path

# Gradio 介面設置，兩個文字框和一個儲存文檔的按鈕
iface = gr.Interface(
    fn=transcribe_audio_and_save,
    inputs=gr.Audio(type="filepath", label="上傳音效檔"),  # 上傳音效檔
    outputs=[
        gr.Textbox(label="原始辨識文字"),  # 第一個框顯示原始辨識
        gr.Textbox(label="加入空格後的文字"),  # 第二個框顯示加入空格的辨識
        gr.File(label="下載辨識結果文檔")  # 提供下載結果的文檔
    ],
    title="Whisper 音效辨識與儲存",
    description="上傳音效檔，Whisper 會自動辨識語音，並分別顯示原始輸出和加入空格的輸出，還可以下載結果文檔。"
)

# 啟動介面
iface.launch(share=True)