import streamlit as st
import torch
import librosa
import numpy as np
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
import io
from datetime import datetime
import gc
import warnings
warnings.filterwarnings('ignore')

# Konfiguracja strony i optymalizacja pamięci
st.set_page_config(
    page_title="Transkrypcja Audio - Polski",
    page_icon="🎤",
    layout="centered" # zmniejszone zużycie miejsca
)

# Optymalizacja torch
torch.backends.cudnn.benchmark = True
if torch.cuda.is_available():
    torch.cuda.empty_cache()

@st.cache_resource(ttl=3600)  # cache wygasa po godzinie
def zaladuj_model():
    """Ładuje model i procesor z cache z obsługą błędów"""
    try:
        nazwa_modelu = "jonatasgrosman/wav2vec2-large-xlsr-53-polish"
        procesor = Wav2Vec2Processor.from_pretrained(nazwa_modelu)
        model = Wav2Vec2ForCTC.from_pretrained(nazwa_modelu)
        
        # Optymalizacja modelu
        if torch.cuda.is_available():
            model = model.to('cuda')
        model.eval()  # tryb ewaluacji
        return procesor, model
    except Exception as e:
        st.error(f"Błąd ładowania modelu: {str(e)}")
        return None, None

@st.cache_data(ttl=300)  # cache na 5 minut dla danych audio
def przetworz_audio(audio_bytes):
    """Wstępne przetwarzanie audio z optymalizacją pamięci"""
    try:
        # Używamy małych fragmentów do przetwarzania
        y, sr = librosa.load(io.BytesIO(audio_bytes), sr=16000, mono=True)
        return y, sr
    except Exception as e:
        st.error(f"Błąd przetwarzania audio: {str(e)}")
        return None, None

def transkrybuj_audio(audio, procesor, model, chunk_length_s=30):
    """Transkrybuje audio w chunks dla optymalizacji pamięci"""
    try:
        # Podziel audio na chunki
        sample_rate = 16000
        chunk_length = chunk_length_s * sample_rate
        chunks = [audio[i:i + chunk_length] for i in range(0, len(audio), chunk_length)]
        
        pelna_transkrypcja = []
        
        # Przetwarzaj każdy chunk osobno
        for chunk in chunks:
            if len(chunk) < 100:  # pomijamy zbyt krótkie chunki
                continue
                
            inputs = procesor(chunk, sampling_rate=sample_rate, return_tensors="pt", padding=True)
            
            if torch.cuda.is_available():
                inputs = inputs.input_values.to('cuda')
            else:
                inputs = inputs.input_values

            with torch.no_grad():
                logits = model(inputs).logits
                predicted_ids = torch.argmax(logits, dim=-1)
                transkrypcja = procesor.batch_decode(predicted_ids)[0]
                pelna_transkrypcja.append(transkrypcja)
            
            # Czyszczenie pamięci
            del inputs, logits, predicted_ids
            torch.cuda.empty_cache() if torch.cuda.is_available() else gc.collect()
        
        return " ".join(pelna_transkrypcja)
    
    except Exception as e:
        st.error(f"Błąd transkrypcji: {str(e)}")
        return ""

def main():
    st.title("🎤 Transkrypcja Audio w Języku Polskim")
    
    # Ładowanie modelu
    procesor, model = zaladuj_model()
    if procesor is None or model is None:
        st.stop()
    
    # Limit rozmiaru pliku (10MB)
    plik_audio = st.file_uploader(
        "Wybierz plik audio (max 30MB)", 
        type=['wav', 'mp3', 'ogg', 'm4a'],
        accept_multiple_files=False
    )
    
    if plik_audio is not None:
        # Sprawdzenie rozmiaru pliku
        if plik_audio.size > 10 * 1024 * 3072:  # 30MB
            st.error("Plik jest zbyt duży. Maksymalny rozmiar to 30MB.")
            st.stop()
            
        st.audio(plik_audio)
        
        if st.button("Rozpocznij transkrypcję", type="primary"):
            progress_bar = st.progress(0)
            status_text = st.empty()
            
            try:
                # Przetwarzanie audio
                status_text.text("Przetwarzanie audio...")
                progress_bar.progress(25)
                audio, sr = przetworz_audio(plik_audio.getvalue())
                
                if audio is None:
                    st.stop()
                
                # Transkrypcja
                status_text.text("Trwa transkrypcja...")
                progress_bar.progress(50)
                transkrypcja = transkrybuj_audio(audio, procesor, model)
                
                # Wyświetlenie wyniku
                progress_bar.progress(100)
                status_text.text("Zakończono!")
                
                if transkrypcja:
                    st.markdown("### Wynik transkrypcji:")
                    st.text_area("", transkrypcja, height=200)
                    
                    # Przycisk pobierania
                    nazwa_pliku = f"transkrypcja_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
                    st.download_button(
                        "📥 Pobierz transkrypcję",
                        transkrypcja.encode('utf-8'),
                        nazwa_pliku,
                        mime="text/plain"
                    )
                
                # Czyszczenie
                del audio
                gc.collect()
                
            except Exception as e:
                st.error(f"Wystąpił nieoczekiwany błąd: {str(e)}")
            finally:
                progress_bar.empty()
                status_text.empty()

    # Informacje
    with st.expander("ℹ️ Informacje o aplikacji"):
        st.markdown("""
        - Model: Wav2Vec2-Large-XLSR-53-Polish
        - Maksymalny rozmiar pliku: 10MB
        - Obsługiwane formaty: WAV, MP3, OGG, M4A
        - Język: Polski
        """)

if __name__ == "__main__":
    main()