from flask import Flask, request, jsonify
from yt_dlp import YoutubeDL
import os
import uuid
import threading
import hashlib

app = Flask(__name__)

# 存储下载文件的临时目录
DOWNLOAD_DIR = 'downloads'
if not os.path.exists(DOWNLOAD_DIR):
    os.makedirs(DOWNLOAD_DIR)

from flask import send_from_directory

@app.route('/downloads/<path:path>')
def serve_downloads(path):
    return send_from_directory(DOWNLOAD_DIR, path)


@app.route('/')
def root():
    return jsonify({'message': 'This is an API service. Please use the appropriate endpoints.'}), 404

@app.route('/get-info', methods=['POST'])
def get_info():
    data = request.json
    url = data.get('url')

    if not url:
        return jsonify({'error': 'URL is required'}), 400

    try:
        ydl_opts = {
            'cookiefile': 'www.youtube.com_cookies.txt'
        }

        with YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=False)
            return jsonify({
                'title': info['title'],
                'thumbnail': info.get('thumbnail'),
                'duration': info.get('duration'),
                'channel': info.get('channel')
            })

    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/download', methods=['POST'])
def download_audio():
    data = request.json
    url = data.get('url')

    if not url:
        return jsonify({'error': 'URL is required'}), 400

    try:
        # 生成唯一的文件名
        unique_id = str(uuid.uuid4())
        file_name = os.path.join(DOWNLOAD_DIR, f'{unique_id}.mp3')

        # Set up yt-dlp options to download only audio
        ydl_opts = {
            'format': 'bestaudio/best',  # Download the best available audio
            'outtmpl': file_name.rsplit('.', 1)[0] + '.%(ext)s',  # Output filename format
            'cookiefile': 'www.youtube.com_cookies.txt',  # Use cookies file if needed
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',  # Extract audio with FFmpeg
                'preferredcodec': 'mp3',  # Convert to MP3
                'preferredquality': '128',  # Set audio quality
            }],
            'noplaylist': True,  # Avoid downloading entire playlists
        }

        with YoutubeDL(ydl_opts) as ydl:
            # Extract video info and download the audio
            ydl.extract_info(url, download=True)

        # 返回可下载的地址
        download_url = f'{request.host_url}{DOWNLOAD_DIR}/{os.path.basename(file_name)}'
        return jsonify({'download_url': download_url})

    except Exception as e:
        return jsonify({'error': str(e)}), 500


 

# In-memory cache for downloaded videos  
download_cache = {}  # url_hash -> file_path  
download_status = {}  # download_id -> status  

def get_url_hash(url):  
    """Generate a hash of the URL for caching"""  
    return hashlib.md5(url.encode()).hexdigest()  

def download_in_background(url, file_name, download_id, url_hash):  
    try:  
        ydl_opts = {  
            'format': 'bestaudio/best',  
            'outtmpl': file_name.rsplit('.', 1)[0] + '.%(ext)s',  
            'cookiefile': 'www.youtube.com_cookies.txt',  
            'postprocessors': [{  
                'key': 'FFmpegExtractAudio',  
                'preferredcodec': 'mp3',  
                'preferredquality': '320',  
            }],  
            'noplaylist': True,  
        }  
        
        with YoutubeDL(ydl_opts) as ydl:  
            ydl.extract_info(url, download=True)  
        
        # Update cache and status  
        actual_file = f"{file_name.rsplit('.', 1)[0]}.mp3"  
        download_cache[url_hash] = actual_file  
        download_status[download_id] = {  
            'status': 'completed',  
            'download_url': f'{request.host_url}{DOWNLOAD_DIR}/{os.path.basename(actual_file)}'  
        }  
    except Exception as e:  
        download_status[download_id] = {  
            'status': 'error',  
            'error': str(e)  
        }  

@app.route('/audio', methods=['POST'])  
def download_high_audio():  
    data = request.json  
    url = data.get('url')  
    
    if not url:  
        return jsonify({'error': 'URL is required'}), 400  
    
    # Check cache first  
    url_hash = get_url_hash(url)  
    cached_file = download_cache.get(url_hash)  
    
    if cached_file and os.path.exists(cached_file):  
        # Return cached file immediately  
        return jsonify({  
            'status': 'completed',  
            'download_url': f'{request.host_url}{DOWNLOAD_DIR}/{os.path.basename(cached_file)}'  
        })  
    
    # Start new download  
    download_id = str(uuid.uuid4())  
    file_name = os.path.join(DOWNLOAD_DIR, f'{download_id}.mp3')  
    
    download_status[download_id] = {'status': 'processing'}  
    thread = Thread(target=download_in_background, args=(url, file_name, download_id, url_hash))  
    thread.daemon = True  
    thread.start()  
    
    return jsonify({  
        'download_id': download_id,  
        'status': 'processing'  
    })  

@app.route('/audio/status/<download_id>', methods=['GET'])  
def check_download_status(download_id):  
    status = download_status.get(download_id, {'status': 'not_found'})  
    return jsonify(status)


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=True)