From f693dd07ffcbccb3942e2325a83e1208c40fd752 Mon Sep 17 00:00:00 2001
From: tavo-wasd <gustavo@gustavocalvo.xyz>
Date: Sun, 7 Apr 2024 23:44:22 -0600
Subject: [PATCH] play from yt

---
 bootstrap.sh |  3 +-
 main.py      | 92 +++++++++++++++++++++++++++++++++++-----------------
 2 files changed, 64 insertions(+), 31 deletions(-)

diff --git a/bootstrap.sh b/bootstrap.sh
index e9b5243..2cff236 100644
--- a/bootstrap.sh
+++ b/bootstrap.sh
@@ -5,6 +5,7 @@
     . ./bin/activate ;
     pip install discord ;
     pip install PyNaCl ;
-    pip install fuzzywuzzy ;
     pip install python-dotenv ;
+    pip install ytmusicapi ;
+    pip install yt_dlp ;
 )
diff --git a/main.py b/main.py
index baaa41e..9cfea75 100644
--- a/main.py
+++ b/main.py
@@ -2,39 +2,65 @@ import discord
 from discord.ext import commands
 import asyncio
 import os
-from fuzzywuzzy import process
 from dotenv import load_dotenv
+import yt_dlp
+from ytmusicapi import YTMusic
 
 load_dotenv()
+songs_dir = os.getenv("SONGS_DIR")
 intents = discord.Intents.default()
 intents.message_content = True
 intents.voice_states = True
 bot = commands.Bot(command_prefix='!', intents=intents)
 queued_songs = []
-voice_client = None  # Global variable to track voice client
+voice_client = None
 
-def find_song(query, songs_dir):
-    best_match = None
-    best_confidence = 0
+def search_song(query):
+    ytmusic = YTMusic()
+    search_results = ytmusic.search(query, filter='songs', limit=1)
+    if search_results:
+        song_id = search_results[0]['videoId']
+        return song_id
+    else:
+        return None
+
+def search_video(query):
+    ydl_opts = {
+        'default_search': 'auto',
+        'format': 'best',
+        'quiet': True
+    }
     
-    for root, dirs, files in os.walk(songs_dir):
-        for file in files:
-            if file.endswith(".mp3"):
-                song_path = os.path.join(root, file)
-                song_name = os.path.splitext(file)[0]
-                full_song_name = os.path.relpath(song_path, start=songs_dir)
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        try:
+            result = ydl.extract_info(query, download=False)
+            if 'entries' in result:
+                # Take the first result
+                video_id = result['entries'][0]['id']
+                return video_id
+            else:
+                return None
+        except yt_dlp.DownloadError:
+            return None
 
-                confidence = process.extractOne(query, [song_name.replace("_", " ")])[1]
-                if confidence > best_confidence:
-                    best_match = song_path
-                    best_confidence = confidence
-
-                confidence = process.extractOne(query, [full_song_name.replace("_", " ")])[1]
-                if confidence > best_confidence:
-                    best_match = song_path
-                    best_confidence = confidence
-    
-    return best_match, best_confidence
+def download_song(query):
+    song_id = search_video(query)
+    ydl_opts = {
+        'format': 'bestaudio/best',
+        'outtmpl': f'{songs_dir}/%(title)s.%(ext)s',
+        'postprocessors': [{
+            'key': 'FFmpegExtractAudio',
+            'preferredcodec': 'mp3',
+            'preferredquality': '192',
+        }],
+    }
+    if not os.path.exists(songs_dir):
+        os.makedirs(songs_dir)
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        info = ydl.extract_info(f'https://music.youtube.com/watch?v={song_id}', download=True)
+        filename = ydl.prepare_filename(info)
+        filename = os.path.splitext(filename)[0] + '.mp3'
+    return filename
 
 @bot.command()
 async def play(ctx, *, query):
@@ -53,6 +79,10 @@ async def skip(ctx):
 @bot.command()
 async def stop(ctx):
     global voice_client
+    for file in os.listdir(songs_dir):
+        file_path = os.path.join(songs_dir, file)
+        if os.path.isfile(file_path):
+            os.unlink(file_path)
     if voice_client and voice_client.is_connected():
         queued_songs.clear()
         if voice_client.is_playing():
@@ -63,11 +93,10 @@ async def stop(ctx):
 
 async def add_to_queue(ctx, query):
     global voice_client
-    songs_dir = os.getenv("SONGS_DIR")
-    # Find the closest match to the query
-    song_path, confidence = find_song(query, songs_dir)
+    song_path = download_song(query)
+    print(song_path)
     
-    if confidence >= 70 and song_path:  # Adjust confidence threshold as needed
+    if song_path:
         song_name = os.path.relpath(song_path, start=songs_dir)  # Get the song name relative to songs_dir
         queued_songs.append((song_path, song_name))
         await ctx.send(f"**Added to queue:** `{song_name}`")
@@ -78,11 +107,10 @@ async def connect_and_play(ctx, query):
     global voice_client
     voice_channel = ctx.author.voice.channel
     if voice_channel:
-        songs_dir = os.getenv("SONGS_DIR")
-        # Find the closest match to the query
-        song_path, confidence = find_song(query, songs_dir)
+        song_path = download_song(query)
+        print(song_path)
         
-        if confidence >= 70 and song_path:  # Adjust confidence threshold as needed
+        if song_path:
             song_name = os.path.relpath(song_path, start=songs_dir)  # Get the song name relative to songs_dir
             queued_songs.append((song_path, song_name))
             voice_client = await play_queued_songs(ctx, voice_channel)
@@ -99,6 +127,10 @@ async def play_queued_songs(ctx, voice_channel):
         while voice_client and voice_client.is_playing():  # Check if voice_client is not None
             await asyncio.sleep(1)
     if voice_client:
+        for file in os.listdir(songs_dir):
+            file_path = os.path.join(songs_dir, file)
+            if os.path.isfile(file_path):
+                os.unlink(file_path)
         await voice_client.disconnect()
 
 bot.run(os.getenv("DISCORD_TOKEN"))