import os
import re
import yt_dlp
from pytube import YouTube

class VideoDownloader:
    def __init__(self, download_path='downloads'):
        self.download_path = download_path
        # Ensure download directory exists
        os.makedirs(download_path, exist_ok=True)
    
    def get_platform(self, url):
        """Determine the platform from the URL"""
        if 'youtube.com' in url or 'youtu.be' in url:
            return 'youtube'
        elif 'facebook.com' in url or 'fb.com' in url or 'fb.watch' in url:
            return 'facebook'
        elif 'instagram.com' in url:
            return 'instagram'
        elif 'tiktok.com' in url:
            return 'tiktok'
        else:
            return 'generic'
    
    def download_youtube(self, url, audio_only=False):
        """Download video from YouTube using pytube"""
        try:
            yt = YouTube(url)
            if audio_only:
                # Download audio only
                stream = yt.streams.filter(only_audio=True).first()
                out_file = stream.download(output_path=self.download_path)
                # Convert to mp3
                base, _ = os.path.splitext(out_file)
                mp3_file = base + '.mp3'
                os.rename(out_file, mp3_file)
                return {'success': True, 'file_path': mp3_file, 'title': yt.title, 'format': 'mp3'}
            else:
                # Download video with highest resolution
                stream = yt.streams.get_highest_resolution()
                out_file = stream.download(output_path=self.download_path)
                return {'success': True, 'file_path': out_file, 'title': yt.title, 'format': 'mp4'}
        except Exception as e:
            return {'success': False, 'error': str(e)}
    
    def download_with_ytdlp(self, url, audio_only=False):
        """Download video using yt-dlp for all platforms"""
        try:
            platform = self.get_platform(url)
            filename = f"{platform}_{os.urandom(4).hex()}"
            output_template = os.path.join(self.download_path, filename)
            
            ydl_opts = {
                'outtmpl': output_template + '.%(ext)s',
                'quiet': True,
                'no_warnings': True,
            }
            
            if audio_only:
                ydl_opts.update({
                    'format': 'bestaudio/best',
                    'postprocessors': [{
                        'key': 'FFmpegExtractAudio',
                        'preferredcodec': 'mp3',
                        'preferredquality': '192',
                    }],
                })
                final_ext = 'mp3'
            else:
                ydl_opts.update({
                    'format': 'best',
                })
                final_ext = 'mp4'  # Most common, but might be different
            
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                if 'entries' in info:  # Playlist
                    info = info['entries'][0]
                
                # Get the actual file extension
                if not audio_only:
                    final_ext = info.get('ext', 'mp4')
                
                title = info.get('title', 'Video')
                file_path = f"{output_template}.{final_ext}"
                
                return {
                    'success': True, 
                    'file_path': file_path, 
                    'title': title, 
                    'format': final_ext
                }
        except Exception as e:
            return {'success': False, 'error': str(e)}
    
    def download(self, url, audio_only=False):
        """Main download method that routes to appropriate downloader"""
        platform = self.get_platform(url)
        
        # Use pytube for YouTube if not audio_only (more reliable for video)
        if platform == 'youtube' and not audio_only:
            return self.download_youtube(url, audio_only)
        else:
            # Use yt-dlp for all other platforms and YouTube audio
            return self.download_with_ytdlp(url, audio_only)
