import os
from flask import Flask, render_template, request, jsonify, send_file
import threading
import uuid
from pathlib import Path
from yt_dlp import YoutubeDL
from PIL import Image

app = Flask(__name__, template_folder=os.getcwd())  # Set template_folder to the current directory

downloads = {}

class MusicDownloader:
    def __init__(self):
        self.download_dir = Path("downloads")
        self.download_dir.mkdir(exist_ok=True)

    def download(self, url: str, download_id: str) -> bool:
        downloads[download_id] = {'status': 'downloading', 'progress': '0%'}
        
ydl_opts = {
    'format': 'bestaudio/best',
    'postprocessors': [{
        'key': 'FFmpegExtractAudio',
        'preferredcodec': 'mp3',
        'preferredquality': '192',
    }],
    'outtmpl': str(self.download_dir / '%(title)s.%(ext)s'),
    'writethumbnail': True,
    'cookiefile': 'cookies.txt',  # Path to your cookies file
    'progress_hooks': [lambda d: self._progress_hook(d, download_id)],
}

        try:
            with YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                title = info.get('title', 'unknown_title')

                webp_path = self.download_dir / f"{title}.webp"
                if webp_path.exists():
                    self._convert_thumbnail(webp_path)

                downloads[download_id]['status'] = 'completed'
                downloads[download_id]['title'] = title
                return True

        except Exception as e:
            downloads[download_id]['status'] = 'failed'
            downloads[download_id]['error'] = str(e)
            return False

    def _progress_hook(self, d, download_id):
        if d['status'] == 'downloading':
            downloads[download_id]['progress'] = d.get('_percent_str', '0%')

    def _convert_thumbnail(self, thumbnail_path: Path):
        try:
            with Image.open(thumbnail_path) as img:
                jpg_path = thumbnail_path.with_suffix('.jpg')
                img.convert('RGB').save(jpg_path, 'JPEG')
                thumbnail_path.unlink()
        except Exception:
            pass

@app.route('/')
def index():
    return render_template('index.html')  # No need for templates folder, index.html is in the same folder

@app.route('/download', methods=['POST'])
def start_download():
    url = request.json.get('url')
    if not url:
        return jsonify({'error': 'No URL provided'}), 400

    download_id = str(uuid.uuid4())
    downloader = MusicDownloader()

    threading.Thread(target=downloader.download, args=(url, download_id)).start()

    return jsonify({'download_id': download_id})

@app.route('/status/<download_id>')
def get_status(download_id):
    if download_id not in downloads:
        return jsonify({'status': 'not_found'}), 404
    return jsonify(downloads[download_id])

@app.route('/download/<download_id>')
def get_file(download_id):
    if download_id not in downloads or downloads[download_id]['status'] != 'completed':
        return 'File not found', 404

    title = downloads[download_id]['title']
    file_path = os.path.join('downloads', f'{title}.mp3')

    if not os.path.exists(file_path):
        return 'File not found', 404

    return send_file(file_path, as_attachment=True)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)  # Port set for Hugging Face Spaces