import streamlit as st
from st_audiorec import st_audiorec
# import whisper
import numpy as np

# model = whisper.load_model("base")

# def transcribe(audio):
    
#     #time.sleep(3)
#     # load audio and pad/trim it to fit 30 seconds
#     audio = whisper.load_audio(audio)
#     audio = whisper.pad_or_trim(audio)

#     # make log-Mel spectrogram and move to the same device as the model
#     mel = whisper.log_mel_spectrogram(audio).to(model.device)

#     # detect the spoken language
#     _, probs = model.detect_language(mel)
#     print(f"Detected language: {max(probs, key=probs.get)}")

#     # decode the audio
#     options = whisper.DecodingOptions(fp16 = False)
#     result = whisper.decode(model, mel, options)
#     return result.text

import torch
from transformers import pipeline

# start model
device = "cuda:0" if torch.cuda.is_available() else "cpu"

translation_pipe = pipeline(
        "automatic-speech-recognition", model="openai/whisper-base", device=device
    )

def translate(audio, lang = "en"):
    outputs = translation_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "transcribe", "language": lang})
    return outputs["text"]
    
    
# Customizing the Streamlit page layout
st.set_page_config(
    page_title="tłumacz",
    page_icon="🎤",
    layout="wide",
)

# Center the app title and adjust padding
st.markdown(
    """
    <style>
    .stApp {
        max-width: 1000px;
        margin: 0 auto;
        padding: 2rem;
    }
    .stTitle {
        text-align: center;
    }
    </style>
    """,
    unsafe_allow_html=True
)

def audiorec_demo_app():

    st.title('Audio Recorder App')
    st.subheader('Record and Playback Audio')

    # Add a little info on how to use the app
    st.markdown(
        "🎤 To record audio, simply click the 'Start Recording' button. "
        "Once you're done, click 'Stop Recording' and 'Save Recording' to play it back. "
        "You can also reset the recording if needed."
    )

    wav_audio_data = st_audiorec()

    filename = "temp.wav"

    if wav_audio_data is not None:
        # Display the recorded audio
        st.subheader('Transcribed text:')
        with st.spinner("Transcribing..."):

            with open(filename, "wb") as f:
                f.write(wav_audio_data)
            with open(filename, "rb") as f:
                audio = f.read()

            output_text = translate(audio, lang="pl")
            translated_text = translate(audio, lang="en")
            st.markdown(f"**Orginalny tekst [PL]:** {output_text}")
            st.markdown(f"**Tłumaczenie [EN]:** {translated_text}")

if __name__ == '__main__':
    audiorec_demo_app()