you-music/backend/app/services/downloader.py

import asyncio
import os
import re
import subprocess
from typing import Optional, Tuple
from pathlib import Path
import aiohttp
import mutagen
from mutagen.mp3 import MP3
from mutagen.id3 import ID3, TIT2, TPE1, TALB, APIC
from app.core.config import settings
import logging

logger = logging.getLogger(__name__)


class MusicDownloader:
    """Download music from various sources using yt-dlp (similar to xiaomusic)"""
    
    def __init__(self):
        self.download_path = settings.MUSIC_DIR
        self.temp_path = settings.TEMP_DIR
        self.proxy = settings.PROXY
        self.ffmpeg_location = settings.FFMPEG_LOCATION
        
    async def download_music(
        self, 
        url: str, 
        output_name: Optional[str] = None
    ) -> Tuple[bool, str, Optional[str]]:
        """
        Download music from URL using yt-dlp
        Returns: (success, file_path, error_message)
        """
        try:
            # First, extract info to get metadata
            info_cmd = [
                "yt-dlp",
                "--dump-json",
                "--no-playlist",
                url
            ]
            
            if self.proxy:
                info_cmd.extend(["--proxy", self.proxy])
            
            info_process = await asyncio.create_subprocess_exec(
                *info_cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            info_stdout, _ = await info_process.communicate()
            
            # Extract metadata from info
            import json
            video_info = {}
            artist_name = None
            thumbnail_url = None
            if info_process.returncode == 0 and info_stdout:
                try:
                    video_info = json.loads(info_stdout.decode())
                    # Try to extract artist from various fields
                    artist_name = (
                        video_info.get('artist') or 
                        video_info.get('creator') or 
                        video_info.get('uploader') or 
                        video_info.get('channel')
                    )
                    # Get best thumbnail
                    thumbnail_url = video_info.get('thumbnail')
                except:
                    pass
            
            # Prepare output template
            if output_name:
                title = f"{output_name}.%(ext)s"
            else:
                title = "%(title)s.%(ext)s"
            
            # Build yt-dlp command arguments (similar to xiaomusic)
            cmd_args = [
                "yt-dlp",
                "--no-playlist",
                "-x",  # Extract audio
                "--audio-format", settings.YT_DLP_AUDIO_FORMAT,
                "--audio-quality", settings.YT_DLP_AUDIO_QUALITY,
                "--paths", self.download_path,
                "-o", title,
                "--ffmpeg-location", self.ffmpeg_location,
                "--embed-metadata",  # Embed metadata
                "--parse-metadata", "%(artist)s:%(meta_artist)s",  # Parse artist
                "--parse-metadata", "%(uploader)s:%(meta_artist)s",  # Fallback to uploader
            ]
            
            if self.proxy:
                cmd_args.extend(["--proxy", self.proxy])
            
            cmd_args.append(url)
            
            logger.info(f"Downloading: {' '.join(cmd_args)}")
            
            # Execute download
            process = await asyncio.create_subprocess_exec(
                *cmd_args,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            stdout, stderr = await process.communicate()
            
            if process.returncode == 0:
                # Find the downloaded file
                output_file = await self._find_downloaded_file(output_name)
                if output_file:
                    # Download and embed thumbnail from online source
                    # Priority: Online thumbnail -> Embedded in metadata extraction
                    thumbnail_path = None
                    if thumbnail_url:
                        thumbnail_path = await self._download_thumbnail(thumbnail_url, output_file)
                        if thumbnail_path:
                            await self._embed_thumbnail(output_file, thumbnail_path)
                    
                    # Embed artist info if we extracted it
                    if artist_name:
                        await self._embed_artist_metadata(output_file, artist_name, video_info)
                    
                    logger.info(f"Download successful: {output_file}")
                    return True, output_file, None
                else:
                    return False, "", "Downloaded file not found"
            else:
                error_msg = stderr.decode() if stderr else "Unknown error"
                logger.error(f"Download failed: {error_msg}")
                return False, "", error_msg
                
        except Exception as e:
            logger.exception(f"Download exception: {e}")
            return False, "", str(e)
    
    async def _find_downloaded_file(self, output_name: Optional[str]) -> Optional[str]:
        """Find the most recently downloaded file"""
        try:
            files = []
            for ext in ['.mp3', '.m4a', '.opus', '.webm']:
                if output_name:
                    pattern = f"{output_name}{ext}"
                    file_path = os.path.join(self.download_path, pattern)
                    if os.path.exists(file_path):
                        return file_path
                else:
                    # Find most recent file
                    for file in Path(self.download_path).glob(f"*{ext}"):
                        files.append(file)
            
            if files:
                # Return most recent file
                latest_file = max(files, key=lambda x: x.stat().st_mtime)
                return str(latest_file)
            
            return None
        except Exception as e:
            logger.error(f"Error finding downloaded file: {e}")
            return None
    
    async def _embed_artist_metadata(self, file_path: str, artist: str, video_info: dict = None):
        """Embed artist and other metadata into MP3 file"""
        try:
            audio = MP3(file_path, ID3=ID3)
            
            # Add ID3 tag if it doesn't exist
            try:
                audio.add_tags()
            except mutagen.id3.error:
                pass
            
            # Set artist
            audio.tags.add(TPE1(encoding=3, text=artist))
            
            # Set title if available
            if video_info and video_info.get('title'):
                audio.tags.add(TIT2(encoding=3, text=video_info['title']))
            
            # Set album if available
            if video_info and video_info.get('album'):
                audio.tags.add(TALB(encoding=3, text=video_info['album']))
            
            audio.save()
            logger.info(f"Embedded artist metadata: {artist}")
        except Exception as e:
            logger.warning(f"Could not embed metadata: {e}")
    
    async def _download_thumbnail(self, thumbnail_url: str, audio_file_path: str) -> Optional[str]:
        """Download thumbnail from URL"""
        try:
            # Create thumbnails directory
            thumbnails_dir = os.path.join(settings.MUSIC_DIR, "thumbnails")
            os.makedirs(thumbnails_dir, exist_ok=True)
            
            # Generate thumbnail filename based on audio file
            audio_basename = os.path.splitext(os.path.basename(audio_file_path))[0]
            thumbnail_path = os.path.join(thumbnails_dir, f"{audio_basename}.jpg")
            
            # Download thumbnail
            async with aiohttp.ClientSession() as session:
                async with session.get(thumbnail_url) as response:
                    if response.status == 200:
                        with open(thumbnail_path, 'wb') as f:
                            f.write(await response.read())
                        logger.info(f"Downloaded thumbnail: {thumbnail_path}")
                        return thumbnail_path
            
            return None
        except Exception as e:
            logger.warning(f"Could not download thumbnail: {e}")
            return None
    
    async def _embed_thumbnail(self, audio_file_path: str, thumbnail_path: str):
        """Embed thumbnail into MP3 file"""
        try:
            audio = MP3(audio_file_path, ID3=ID3)
            
            # Add ID3 tag if it doesn't exist
            try:
                audio.add_tags()
            except mutagen.id3.error:
                pass
            
            # Read thumbnail data
            with open(thumbnail_path, 'rb') as img_file:
                img_data = img_file.read()
            
            # Add cover art
            audio.tags.add(
                APIC(
                    encoding=3,
                    mime='image/jpeg',
                    type=3,  # Cover (front)
                    desc='Cover',
                    data=img_data
                )
            )
            
            audio.save()
            logger.info(f"Embedded thumbnail into {audio_file_path}")
        except Exception as e:
            logger.warning(f"Could not embed thumbnail: {e}")
    
    async def download_playlist(
        self, 
        url: str, 
        playlist_name: str
    ) -> Tuple[bool, list[str], Optional[str]]:
        """Download entire playlist"""
        try:
            output_dir = os.path.join(self.download_path, playlist_name)
            os.makedirs(output_dir, exist_ok=True)
            
            title = f"{playlist_name}/%(title)s.%(ext)s"
            
            cmd_args = [
                "yt-dlp",
                "--yes-playlist",
                "-x",
                "--audio-format", settings.YT_DLP_AUDIO_FORMAT,
                "--audio-quality", settings.YT_DLP_AUDIO_QUALITY,
                "--paths", self.download_path,
                "-o", title,
                "--ffmpeg-location", self.ffmpeg_location,
            ]
            
            if self.proxy:
                cmd_args.extend(["--proxy", self.proxy])
            
            cmd_args.append(url)
            
            logger.info(f"Downloading playlist: {' '.join(cmd_args)}")
            
            process = await asyncio.create_subprocess_exec(
                *cmd_args,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            await process.wait()
            
            # Get downloaded files
            downloaded_files = []
            if os.path.exists(output_dir):
                for file in Path(output_dir).glob("*.mp3"):
                    downloaded_files.append(str(file))
            
            return True, downloaded_files, None
            
        except Exception as e:
            logger.exception(f"Playlist download exception: {e}")
            return False, [], str(e)
    
    async def get_music_metadata(self, file_path: str) -> dict:
        """Extract metadata from audio file using mutagen
        
        Thumbnail Priority:
        1. Embedded thumbnail in audio file (already downloaded from online during download)
        2. No online fetch here - thumbnails are embedded during download phase
        """
        try:
            # First get basic metadata with easy=True
            audio_easy = mutagen.File(file_path, easy=True)
            if audio_easy is None:
                return {}
            
            metadata = {
                "title": audio_easy.get("title", [os.path.basename(file_path)])[0] if audio_easy.get("title") else os.path.basename(file_path),
                "artist": audio_easy.get("artist", ["Unknown"])[0] if audio_easy.get("artist") else "Unknown",
                "album": audio_easy.get("album", [""])[0] if audio_easy.get("album") else "",
                "duration": audio_easy.info.length if hasattr(audio_easy, 'info') else 0,
            }
            
            # Now extract thumbnail from ID3 tags
            try:
                audio = MP3(file_path, ID3=ID3)
                if audio.tags:
                    # Look for APIC (attached picture) frames
                    for tag in audio.tags.values():
                        if isinstance(tag, APIC):
                            # Save thumbnail to thumbnails directory
                            thumbnails_dir = os.path.join(settings.MUSIC_DIR, "thumbnails")
                            os.makedirs(thumbnails_dir, exist_ok=True)
                            
                            audio_basename = os.path.splitext(os.path.basename(file_path))[0]
                            thumbnail_filename = f"{audio_basename}.jpg"
                            thumbnail_path = os.path.join(thumbnails_dir, thumbnail_filename)
                            
                            # Save thumbnail
                            with open(thumbnail_path, 'wb') as img_file:
                                img_file.write(tag.data)
                            
                            # Store relative path for database
                            metadata["thumbnail"] = f"thumbnails/{thumbnail_filename}"
                            break
            except Exception as e:
                logger.warning(f"Could not extract thumbnail: {e}")
            
            return metadata
        except Exception as e:
            logger.error(f"Error extracting metadata: {e}")
            return {
                "title": os.path.basename(file_path),
                "artist": "Unknown",
                "album": "",
                "duration": 0,
            }
    
    async def get_duration(self, file_path: str) -> float:
        """Get audio duration"""
        try:
            audio = mutagen.File(file_path)
            if audio and hasattr(audio, 'info'):
                return audio.info.length
            return 0.0
        except Exception:
            return 0.0
    
    def extract_youtube_id(self, url: str) -> Optional[str]:
        """Extract YouTube video ID from URL"""
        patterns = [
            r'(?:youtube\.com\/watch\?v=|youtu\.be\/)([a-zA-Z0-9_-]{11})',
            r'youtube\.com\/embed\/([a-zA-Z0-9_-]{11})',
        ]
        
        for pattern in patterns:
            match = re.search(pattern, url)
            if match:
                return match.group(1)
        
        return None
    
    def is_youtube_url(self, url: str) -> bool:
        """Check if URL is from YouTube"""
        return 'youtube.com' in url or 'youtu.be' in url
    
    def is_bilibili_url(self, url: str) -> bool:
        """Check if URL is from Bilibili"""
        return 'bilibili.com' in url


# Global instance
music_downloader = MusicDownloader()