import os
import subprocess
import yt_dlp

def download_audio_from_youtube(url_or_query: str, output_dir: str = "temp_audio") -> str:
    """
    Downloads audio from YouTube using yt-dlp and converts it to a mono WAV file at 16kHz.
    Accepts a YouTube URL or a search query (e.g., "ytsearch1:Hillsong Hosanna").
    Returns: Path to the downloaded .wav file, or None if extraction fails.
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Output file template (using a unique identifier)
    # We will generate a unique filename
    output_filename = os.path.join(output_dir, "%(id)s.%(ext)s")
    
    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': output_filename,
        # Convert to WAV
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'wav',
            'preferredquality': '192',
        }],
        # Limit search results to 1 if it is a query
        'noplaylist': True,
        'quiet': True,
        'no_warnings': True,
    }
    
    # Check if the input is a YouTube URL, if not, treat it as a search query
    is_url = url_or_query.startswith("http://") or url_or_query.startswith("https://")
    if not is_url:
        search_query = f"ytsearch1:{url_or_query} audio"
    else:
        search_query = url_or_query

    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(search_query, download=True)
            if 'entries' in info:
                # Search query returns a list under entries
                video_info = info['entries'][0]
            else:
                video_info = info
                
            video_id = video_info['id']
            downloaded_file = os.path.join(output_dir, f"{video_id}.wav")
            
            # Postprocess to 16kHz mono WAV for optimal Whisper / Librosa usage
            resampled_file = os.path.join(output_dir, f"{video_id}_16k.wav")
            
            # We use ffmpeg to convert to 16kHz mono
            ffmpeg_cmd = [
                'ffmpeg', '-y',
                '-i', downloaded_file,
                '-ar', '16000',
                '-ac', '1',
                resampled_file
            ]
            
            # Run ffmpeg
            subprocess.run(ffmpeg_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            
            # Delete original WAV file to save space
            if os.path.exists(downloaded_file):
                os.remove(downloaded_file)
                
            return resampled_file
            
    except Exception as e:
        print(f"Error downloading audio from YouTube: {e}")
        return None

if __name__ == "__main__":
    # Test query
    path = download_audio_from_youtube("Hosanna Hillsong")
    print(f"Downloaded to: {path}")