import streamlit as st
from streamlit_tools.tools import load_audio_file
from openai_models.whisper import whisper_os
import time
from documents.docs import generate_docx
import os


pipe = whisper_os()


def main(pipe):
    audio_full_path: str = load_audio_file()
    if "validador" not in st.session_state.keys():
        st.session_state.validador = True

    if st.session_state.validador == True and audio_full_path:
        with st.spinner("Transcribiendo..."):
            transcription = pipe(audio_full_path, return_timestamps=True,
                                 generate_kwargs={"language": "spanish"})

            def transcription_generator():
                for word in transcription["text"].split(' '):
                    time.sleep(0.2)
                    yield word + ' '
            st.write_stream(transcription_generator(), )
            if "transcription" not in st.session_state.keys():
                st.session_state.transcription = transcription["text"]
            # st.write(transcription)

            generate_docx(transcription["text"])
            st.session_state.validador = False

    if audio_full_path and os.path.exists(st.session_state.full_path_docx):
        with open(st.session_state.full_path_docx, "rb") as file:
            btn = st.download_button(
                label="Download docx",
                data=file,
                file_name=st.session_state.full_path_docx,
            )
            if btn:
                st.write(st.session_state.transcription)


if __name__ == "__main__":
    main(pipe)