From a87385db78989e9d894019d5a72c65fbaa8fdef7 Mon Sep 17 00:00:00 2001
From: Alexander-D-Karpov <alexandr.d.karpov@gmail.com>
Date: Thu, 1 Feb 2024 15:41:12 +0300
Subject: [PATCH] Implemented URL parsing and music identification improvements

---
 akarpov/music/services/spotify.py | 37 +++++++++++++++---------
 akarpov/music/tasks.py            | 47 ++++++++++++++++++++++++-------
 2 files changed, 60 insertions(+), 24 deletions(-)

diff --git a/akarpov/music/services/spotify.py b/akarpov/music/services/spotify.py
index 3abb3e6..3c623a6 100644
--- a/akarpov/music/services/spotify.py
+++ b/akarpov/music/services/spotify.py
@@ -1,3 +1,5 @@
+import threading
+
 import spotipy
 from django.conf import settings
 from spotdl import Song, Spotdl
@@ -23,21 +25,28 @@ def search(name: str, session: spotipy.Spotify, search_type="track"):
     return res
 
 
+thread_local = threading.local()
+
+
+def get_spotdl_client():
+    if not hasattr(thread_local, "spotdl_client"):
+        spot_settings = {
+            "simple_tui": True,
+            "log_level": "ERROR",
+            "lyrics_providers": ["genius", "musixmatch"],
+        }
+        thread_local.spotdl_client = Spotdl(
+            client_id=settings.MUSIC_SPOTIFY_ID,
+            client_secret=settings.MUSIC_SPOTIFY_SECRET,
+            user_auth=False,
+            headless=False,
+            downloader_settings=spot_settings,
+        )
+    return thread_local.spotdl_client
+
+
 def download_url(url, user_id=None):
-    spot_settings = {
-        "simple_tui": True,
-        "log_level": "ERROR",
-        "lyrics_providers": ["genius", "musixmatch"],
-    }
-
-    spotdl_client = Spotdl(
-        client_id=settings.MUSIC_SPOTIFY_ID,
-        client_secret=settings.MUSIC_SPOTIFY_SECRET,
-        user_auth=False,
-        headless=False,
-        downloader_settings=spot_settings,
-    )
-
+    spotdl_client = get_spotdl_client()
     session = create_session()
 
     if "track" in url:
diff --git a/akarpov/music/tasks.py b/akarpov/music/tasks.py
index 8c2a1da..2114163 100644
--- a/akarpov/music/tasks.py
+++ b/akarpov/music/tasks.py
@@ -1,7 +1,9 @@
 from datetime import timedelta
 
 import pylast
+import spotipy
 import structlog
+import yt_dlp
 from asgiref.sync import async_to_sync
 from celery import shared_task
 from channels.layers import get_channel_layer
@@ -9,6 +11,7 @@
 from django.utils import timezone
 from django.utils.timezone import now
 from pytube import Channel, Playlist
+from spotipy import SpotifyClientCredentials
 
 from akarpov.music.api.serializers import SongSerializer
 from akarpov.music.models import (
@@ -28,22 +31,46 @@
 
 @shared_task
 def list_tracks(url, user_id):
-    if "music.youtube.com" in url:
+    if "music.youtube.com" in url or "youtu.be" in url:
         url = url.replace("music.youtube.com", "youtube.com")
+        url = url.replace("youtu.be", "youtube.com")
     if "spotify.com" in url:
         spotify.download_url(url, user_id)
     elif "music.yandex.ru" in url:
         yandex.load_playlist(url, user_id)
-    elif "channel" in url or "/c/" in url:
-        p = Channel(url)
-        for video in p.video_urls:
-            process_yb.apply_async(kwargs={"url": video, "user_id": user_id})
-    elif "playlist" in url or "&list=" in url:
-        p = Playlist(url)
-        for video in p.video_urls:
-            process_yb.apply_async(kwargs={"url": video, "user_id": user_id})
+    if "youtube.com" in url:
+        if "channel" in url or "/c/" in url:
+            channel = Channel(url)
+            for video in channel.videos:
+                with yt_dlp.YoutubeDL({}) as ydl:
+                    info = ydl.extract_info(video, download=False)
+                    if info.get("category") == "Music":
+                        process_yb.apply_async(
+                            kwargs={"url": video.watch_url, "user_id": user_id}
+                        )
+        elif "playlist" in url or "&list=" in url:
+            playlist = Playlist(url)
+            for video in playlist.videos:
+                process_yb.apply_async(
+                    kwargs={"url": video.watch_url, "user_id": user_id}
+                )
+        else:
+            process_yb.apply_async(kwargs={"url": url, "user_id": user_id})
     else:
-        process_yb.apply_async(kwargs={"url": url, "user_id": user_id})
+        spotify_manager = SpotifyClientCredentials(
+            client_id=settings.MUSIC_SPOTIFY_ID,
+            client_secret=settings.MUSIC_SPOTIFY_SECRET,
+        )
+        spotify_search = spotipy.Spotify(client_credentials_manager=spotify_manager)
+
+        results = spotify_search.search(q=url, type="track", limit=1)
+        top_track = (
+            results["tracks"]["items"][0] if results["tracks"]["items"] else None
+        )
+
+        if top_track:
+            spotify.download_url(top_track["external_urls"]["spotify"], user_id)
+            url = top_track["external_urls"]["spotify"]
 
     return url