import streamlit as st
import tempfile
from utils.noise_removal import remove_noise
from utils.vad_segmentation import vad_segmentation
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')

    try:
        st.subheader("1️⃣ Noise Removal")
        denoised_path = tmp_path.replace(".wav", "_denoised.wav")
        with st.spinner("Removing noise..."):
            remove_noise(tmp_path, denoised_path)
        st.audio(denoised_path, format="audio/wav")
    except Exception as e:
        st.error(f"Noise removal failed: {e}")

    try:
        st.subheader("2️⃣ Speech Segmentation")
        with st.spinner("Running Voice Activity Detection..."):
            speech_annotation = vad_segmentation(denoised_path)
            segments = [(seg.start, seg.end) for seg in speech_annotation.itersegments()]
        st.write(f" Detected {len(segments)} speech segments.")
        for i, (start, end) in enumerate(segments[:5]):
            st.write(f"Segment {i+1}: {start:.2f}s to {end:.2f}s")
    except Exception as e:
        st.error(f"VAD failed: {e}")

    try:
        st.subheader("3️⃣ Speaker Diarization")
        with st.spinner("Diarizing speakers..."):
            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 {speaker}")
    except Exception as e:
        st.error(f"Speaker diarization failed: {e}")

    try:
        st.subheader("4️⃣ Noise Classification")
        with st.spinner("Classifying background noise..."):
            noise_predictions = classify_noise(denoised_path)
        st.write(" Top predicted noise classes:")
        for cls, prob in noise_predictions:
            st.write(f"{cls}: {prob:.2f}")
    except Exception as e:
        st.error(f"Noise classification failed: {e}")