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

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def download_audio(url):
    try:
        # Configure YouTube client
        logger.info("Initializing YouTube object...")
        yt = YouTube(url, 'WEB')
            
        # 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 memory buffer
        logger.info("Downloading audio...")
        buffer = BytesIO()
        audio_stream.stream_to_buffer(buffer)
        buffer.seek(0)
        
        # Create temp file with correct extension
        file_ext = audio_stream.mime_type.split('/')[-1]
        with tempfile.NamedTemporaryFile(suffix=f".{file_ext}", delete=False) as tmp_file:
            tmp_file.write(buffer.read())
            tmp_path = tmp_file.name
            
        # Generate HTML5 audio player
        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 style="width: 100%">
            <source src="data:audio/{file_ext};base64,{audio_base64}" type="audio/{file_ext}">
            Your browser does not support the audio element.
        </audio>
        """
        
        return [
            f"Success: {yt.title}",
            tmp_path,
            audio_html
        ]
        
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        return [f"Error: {str(e)}", None, None]

# Gradio interface
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...",
            max_lines=1
        )
        
        with gr.Column(scale=0.3):
            clear_btn = gr.Button("Clear", variant="secondary")
            submit_btn = gr.Button("Submit", variant="primary")
    
    status_output = gr.Textbox(label="Status", interactive=False)
    file_download = gr.File(label="Downloaded Audio")
    audio_player = gr.HTML(label="Audio Preview")
    
    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=["/"]
    )