import streamlit as st
import librosa
import soundfile as sf
import tempfile
import os
from utils.noise_removal import remove_noise
from utils.vad_segmentation import detect_speech_segments
from utils.speaker_diarization import diarize_speakers
from utils.noise_classification import classify_noise

st.set_page_config(page_title="Audio Analyzer", layout="wide")
st.title(" Audio Analysis Pipeline")

uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "mp3"])

if uploaded_file:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        tmp.write(uploaded_file.read())
        tmp_path = tmp.name

    st.audio(tmp_path, format='audio/wav')

    st.subheader("1️⃣ Noise Removal")
    denoised_path = tmp_path.replace(".wav", "_denoised.wav")
    remove_noise(tmp_path, denoised_path)
    st.audio(denoised_path, format="audio/wav")

    st.subheader("2️⃣ Speech Segmentation")
    speech_segments = detect_speech_segments(denoised_path)
    st.write(f"Detected {len(speech_segments)} speech segments.")
    for i, (start, end) in enumerate(speech_segments[:5]):
        st.write(f"Segment {i+1}: {start:.2f}s to {end:.2f}s")

    st.subheader("3️⃣ Speaker Diarization")
    diarization = diarize_speakers(denoised_path)
    st.text("Speakers detected:")
    for turn, _, speaker in diarization.itertracks(yield_label=True):
        st.write(f"{turn.start:.2f}s - {turn.end:.2f}s: {speaker}")

    st.subheader("4️⃣ Noise Classification")
    noise_predictions = classify_noise(denoised_path)
    st.write("Top predicted noise classes:")
    for cls, prob in noise_predictions:
        st.write(f"{cls}: {prob:.2f}")