import streamlit as st
from moviepy.editor import VideoFileClip
from speechbrain.pretrained import EncoderClassifier
import torchaudio
import os
from pytube import YouTube

CLASSIFIER = "Jzuluaga/accent-id-commonaccent_xlsr-en-english"

def download_video(url):
    """Handles YouTube and direct video links"""
    if "youtube.com" in url:
        yt = YouTube(url)
        stream = yt.streams.filter(progressive=True, file_extension='mp4').first()
        return stream.download()
    else:  # Direct download
        # Implement custom download logic for non-YouTube links
        return url  

def extract_audio(video_path):
    clip = VideoFileClip(video_path)
    audio_path = "temp_audio.wav"
    clip.audio.write_audiofile(audio_path)
    return audio_path

def classify_accent(audio_path):
    classifier = EncoderClassifier.from_hparams(
        source=CLASSIFIER,
        savedir="pretrained_models/accent_classifier"
    )
    waveform, sample_rate = torchaudio.load(audio_path)
    prediction = classifier.classify_batch(waveform)
    return prediction

def main():
    st.title("REM Waste Accent Classifier")
    url = st.text_input("Enter video URL:", "")
    
    if url:
        with st.spinner("Processing..."):
            video_path = download_video(url)
            audio_path = extract_audio(video_path)
            prediction = classify_accent(audio_path)
            
            st.success("Analysis Complete")
            st.metric("Detected Accent", prediction[3][0])
            st.metric("Confidence Score", f"{prediction[1].exp().max().item()*100:.2f}%")

if __name__ == "__main__":
    main()