import os
import sys
import datetime
import streamlit as st
from asr import load_model, inference
from audiorecorder import audiorecorder

@st.cache_resource
def load_asr_model():
    return load_model()

processor, asr_model = load_asr_model()


def save_audio_file(audio_bytes, file_extension):
    """
    Save audio bytes to a file with the specified extension.

    :param audio_bytes: Audio data in bytes
    :param file_extension: The extension of the output audio file
    :return: The name of the saved audio file
    """
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    file_name = f"audio_{timestamp}.{file_extension}"

    with open(file_name, "wb") as f:
        f.write(audio_bytes)

    return file_name


def transcribe_audio(file_path):
    """
    Transcribe the audio file at the specified path.

    :param file_path: The path of the audio file to transcribe
    :return: The transcribed text
    """
    with open(file_path, "rb") as audio_file:
        transcript = inference(processor, asr_model, audio_file)
    return transcript


def main():
    """
    """
    st.title("Anishinaabemowin Transcription")
    tab1, tab2 = st.tabs(["Record Audio", "Upload Audio"])

    # Record Audio tab
    with tab1:
        audio = audiorecorder("Click to start recording", "Click to stop recording")
        if len(audio) > 0:
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            fname = f"audio_{timestamp}.wav"
            audio.export(fname, format="wav")
            # st.audio(audio_bytes, format="audio/wav")
            # fname = save_audio_file(audio_bytes, "wav")

    # Upload Audio tab
    with tab2:
        audio_file = st.file_uploader("Upload Audio", type=["wav"])
        if audio_file:
            file_extension = audio_file.type.split('/')[1]
            fname = save_audio_file(audio_file.read(), file_extension)

    # Transcribe button action
    if st.button("Click to transcribe uploaded audio."):
        # Transcribe the audio file
        transcript_text = transcribe_audio(fname)

        # Display the transcript
        st.header("Transcript")
        st.write(transcript_text)

        # Save the transcript to a text file
        with open("transcript.txt", "w") as f:
            f.write(transcript_text)

        # Provide a download button for the transcript
        st.download_button("Download Transcript", transcript_text)


if __name__ == "__main__":
    # Set up the working directory
    working_dir = os.path.dirname(os.path.abspath(__file__))
    sys.path.append(working_dir)

    # Run the main function
    main()