import gradio as gr
from pytubefix import YouTube
import tempfile
import base64
from io import BytesIO

def download_audio(url):
    try:
        # Initialize YouTube object
        yt = YouTube(url, 'WEB')
        title = yt.title
        
        # Get best audio stream
        audio_stream = yt.streams.filter(only_audio=True).order_by('abr').desc().first()
        if not audio_stream:
            return ["Error: No audio stream found", None, None]

        # Download to in-memory buffer
        buffer = BytesIO()
        audio_stream.stream_to_buffer(buffer)
        buffer.seek(0)
        
        # Create temporary file
        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp_file:
            tmp_file.write(buffer.read())
            tmp_path = tmp_file.name

        # Create HTML audio player with base64 encoding
        with open(tmp_path, "rb") as f:
            audio_bytes = f.read()
        audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
        audio_html = f"""
        <audio controls autoplay style="width: 100%">
            <source src="data:audio/mp3;base64,{audio_base64}" type="audio/mp3">
            Your browser does not support the audio element.
        </audio>
        """

        return [
            f"Success: {title}",
            tmp_path,
            audio_html
        ]
    
    except Exception as e:
        return [f"Error: {str(e)}", None, None]

with gr.Blocks(title="YouTube Audio Downloader") as demo:
    gr.Markdown("# YouTube Audio Downloader")
    
    with gr.Row():
        url_input = gr.Textbox(label="YouTube URL", placeholder="Enter YouTube URL here...")
        
        with gr.Column(scale=0.3):
            clear_btn = gr.Button("Clear")
            submit_btn = gr.Button("Submit")
    
    status_output = gr.Textbox(label="Status", interactive=False)
    file_download = gr.File(label="Audio File")
    audio_player = gr.HTML(label="Audio Player")
    
    def clear():
        return [None, None, None, None]
    
    clear_btn.click(
        fn=clear,
        outputs=[url_input, status_output, file_download, audio_player]
    )
    
    submit_btn.click(
        fn=download_audio,
        inputs=url_input,
        outputs=[status_output, file_download, audio_player]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        allowed_paths=["*"]
    )