import os
import streamlit as st
from faster_whisper import WhisperModel
from datetime import timedelta
import time

MODEL_PATH = "/content/model"
AUDIO_EXTENSIONS = [".wav", ".mp3", ".flac", ".m4a", ".ogg"]

def load_model():
    return WhisperModel(MODEL_PATH, device="cuda", compute_type="float16")

def get_proces_time(start_time): #時間変換
    end_time = time.time()
    elapsed_time = end_time - start_time
    minutes, seconds = divmod(elapsed_time, 60)
    formatted_time = f"処理時間：{int(minutes)}分{seconds:.1f}秒"
    return formatted_time

def format_segment_tmie(seconds): #音声認識時間のフォーマット関数
    td = timedelta(seconds=seconds)
    return f"{td.seconds // 3600:02}:{(td.seconds // 60) % 60:02}:{td.seconds % 60:02},{td.microseconds // 1000:03}"


def transcribe(audio_file):    
    st.write("モデルを読み込んでいます...") 
    model = load_model()
    file_extension = os.path.splitext(audio_file.name)[1].lower()

    if file_extension in AUDIO_EXTENSIONS:
        segments, info = model.transcribe(audio_file, beam_size=7, language="ja")
        start_time = time.time()
        st.write("処理開始")
        transcribed_text = ""
        for segment in segments:
            output_line = f"[{format_segment_tmie(segment.start)} --> {format_segment_tmie(segment.end)}]{segment.text}"
            st.write(output_line)
            transcribed_text += segment.text + "\n"
            
        formatted_time = get_proces_time(start_time)
        st.write(formatted_time)

        return transcribed_text
    else:
        st.error("Unsupported file format.")
        return ""

st.title("音声ファイルをテキストに変換")

uploaded_file = st.file_uploader("音声ファイルをドロップしてください", type=AUDIO_EXTENSIONS)

if uploaded_file is not None:
    st.audio(uploaded_file)

    
    if st.button("開始"):
        with st.spinner("音声をテキストに変換しています..."):
            transcribed_text = transcribe(uploaded_file)
            st.download_button("テキストファイルをダウンロード", data=transcribed_text.encode(), file_name="transcription.txt", mime="text/plain")
            st.write(transcribed_text)
            
    if st.button("クリア"):
        uploaded_file = None