import gradio as gr
from detect import SimpleOfflineAccentClassifier
import yt_dlp
import os

def download_youtube_audio(url):
    try:
        ydl_opts = {
            'format': 'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'wav',
            }],
            'outtmpl': 'temp_audio',
            'quiet': True,
            'no_warnings': True
        }
        
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
        
        return 'temp_audio.wav'
    except Exception as e:
        return None

def analyze_audio(audio_file, youtube_url):
    if youtube_url:
        audio_file = download_youtube_audio(youtube_url)
        if not audio_file:
            return "Failed to download YouTube audio."
    
    if not audio_file:
        return "Please upload an audio file or provide a YouTube URL."
    
    try:
        classifier = SimpleOfflineAccentClassifier()
        result = classifier.predict_accent(audio_file)
        
        if result is None:
            return "Audio file processing failed."
        
        output = f"Predicted Accent: {result['accent']}\n"
        output += f"Confidence Score: {result['confidence']:.2%}\n\n"
        output += "All Probabilities:\n"
        
        sorted_probs = sorted(
            result['all_probabilities'].items(), 
            key=lambda x: x[1], 
            reverse=True
        )
        
        for accent, prob in sorted_probs:
            bar = "█" * int(prob * 20)
            output += f"- {accent}: {prob:.2%} {bar}\n"
        
        return output
        
    except Exception as e:
        return f"Error occurred: {str(e)}"

# Create Gradio interface
with gr.Blocks() as interface:
    gr.Markdown("# AI Accent Classifier")
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(
                label="Upload Audio File",
                type="filepath"
            )
            
            youtube_url = gr.Textbox(
                label="Or enter YouTube URL",
                placeholder="https://www.youtube.com/watch?v=..."
            )
            
            classify_btn = gr.Button("Analyze Accent")
        
        with gr.Column():
            output_text = gr.Markdown(
                label="Analysis Results",
                value="Analysis results will appear here..."
            )
    
    classify_btn.click(
        fn=analyze_audio,
        inputs=[audio_input, youtube_url],
        outputs=output_text
    )

# Launch the interface
interface.launch()