from flask import Flask, Response, request, stream_with_context
import yt_dlp
import subprocess
import re

app = Flask(__name__)

@app.route('/')
def index():
    return '''
        <form action="/download" method="get">
            <input name="url" placeholder="YouTube URL" required>
            <select name="format">
                <option value="mp4">MP4</option>
                <option value="mp3">MP3</option>
            </select>
            <button type="submit">Download</button>
        </form>
    '''

@app.route('/download')
def download():
    url = request.args.get('url')
    fmt = request.args.get('format', 'mp4')

    # Extract video title first
    with yt_dlp.YoutubeDL({'quiet': True}) as ydl:
        info = ydl.extract_info(url, download=False)
        title = re.sub(r'[^\w\-_\. ]', '_', info['title'])
        filename = f"{title}.{fmt}"

    # Define yt-dlp command to stream output to stdout
    ytdlp_cmd = [
        "yt-dlp",
        url,
        "-f", "bestaudio/best" if fmt == "mp3" else "best",
        "-o", "-",  # output to stdout
    ]

    if fmt == "mp3":
        ytdlp_cmd += [
            "--extract-audio",
            "--audio-format", "mp3",
            "--audio-quality", "192K",
        ]

    def generate():
        process = subprocess.Popen(ytdlp_cmd, stdout=subprocess.PIPE)
        while True:
            chunk = process.stdout.read(1024 * 256)
            if not chunk:
                break
            yield chunk

    return Response(
        stream_with_context(generate()),
        mimetype='application/octet-stream',
        headers={"Content-Disposition": f'attachment; filename="{filename}"'}
    )

if __name__ == '__main__':
    app.run(debug=True)