import streamlit as st
import torch
import librosa
import numpy as np
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
import io
from datetime import datetime

# Konfiguracja strony
st.set_page_config(
    page_title="Transkrypcja Audio - Polski",
    page_icon="🎤",
    layout="wide"
)

@st.cache_resource
def zaladuj_model():
    """Ładuje model i procesor z cache"""
    nazwa_modelu = "jonatasgrosman/wav2vec2-large-xlsr-53-polish"
    procesor = Wav2Vec2Processor.from_pretrained(nazwa_modelu)
    model = Wav2Vec2ForCTC.from_pretrained(nazwa_modelu)
    return procesor, model

def transkrybuj_audio(audio_bytes, procesor, model):
    """Transkrybuje audio z przekazanych bajtów"""
    # Konwersja bajtów na numpy array
    audio, czestotliwosc = librosa.load(io.BytesIO(audio_bytes), sr=16000)
    
    # Przygotowanie danych wejściowych
    dane_wejsciowe = procesor(audio, sampling_rate=16000, return_tensors="pt").input_values

    # Wykonanie transkrypcji
    with torch.no_grad():
        logity = model(dane_wejsciowe).logits
    
    # Dekodowanie transkrypcji
    przewidziane_id = torch.argmax(logity, dim=-1)
    transkrypcja = procesor.batch_decode(przewidziane_id)[0]
    
    return transkrypcja

def main():
    st.title("🎤 Transkrypcja Audio w Języku Polskim")
    
    st.markdown("""
    ### Instrukcja:
    1. Wgraj plik audio (WAV, MP3, etc.)
    2. Poczekaj na transkrypcję
    3. Pobierz wynik jako plik tekstowy
    """)

    # Ładowanie modelu
    with st.spinner("Ładowanie modelu..."):
        procesor, model = zaladuj_model()

    # Upload pliku
    plik_audio = st.file_uploader("Wybierz plik audio", type=['wav', 'mp3', 'ogg', 'm4a'])
    
    if plik_audio is not None:
        st.audio(plik_audio)
        
        if st.button("Rozpocznij transkrypcję"):
            with st.spinner("Trwa transkrypcja..."):
                try:
                    # Transkrypcja
                    transkrypcja = transkrybuj_audio(plik_audio.getvalue(), procesor, model)
                    
                    # Wyświetlenie wyniku
                    st.success("Transkrypcja zakończona!")
                    st.markdown("### Wynik transkrypcji:")
                    st.text_area("", transkrypcja, height=200)
                    
                    # Przygotowanie pliku do pobrania
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    nazwa_pliku = f"transkrypcja_{timestamp}.txt"
                    
                    st.download_button(
                        label="Pobierz transkrypcję",
                        data=transkrypcja.encode('utf-8'),
                        file_name=nazwa_pliku,
                        mime="text/plain"
                    )
                    
                except Exception as e:
                    st.error(f"Wystąpił błąd podczas transkrypcji: {str(e)}")

    st.markdown("---")
    st.markdown("""
    ### Informacje:
    - Model: Wav2Vec2-Large-XLSR-53-Polish
    - Obsługiwane formaty: WAV, MP3, OGG, M4A
    - Język: Polski
    """)

if __name__ == "__main__":
    main()