dj-embe/usecase/sources.py

import json
import re

from discord import Interaction

from entity import Album, Artist, Entry, Playlist, Title
from framework import Spotify, Youtube
from service import FileManager


class Sources:
    def __init__(
        self, fileManager: FileManager, youtube: Youtube, spotify: Spotify
    ) -> None:
        self.fileManager = fileManager

        # youtube
        self.youtube = youtube
        self.youtube_video_regex = re.compile(
            r"(https:\/\/)?(www|music)\.youtube\.com\/(watch\?v=|shorts\/)\w*"
        )
        self.youtube_video_short_regex = re.compile(r"(https:\/\/)?youtu\.be\/\w*")
        self.youtube_playlist_regex = re.compile(
            r"(https:\/\/)?www\.youtube\.com\/playlist\?list=\w*"
        )

        # spotify
        self.spotify = spotify
        self.spotify_track_regex = re.compile(
            r"(https:\/\/)?(open.spotify.com\/|spotify:)track(\/|:)\w*"
        )
        self.spotify_playlist_regex = re.compile(
            r"(https:\/\/)?(open.spotify.com\/|spotify:)playlist(\/|:)\w*"
        )
        self.spotify_album_regex = re.compile(
            r"(https:\/\/)?(open.spotify.com\/|spotify:)album(\/|:)\w*"
        )
        self.spotify_artist_regex = re.compile(
            r"(https:\/\/)?(open.spotify.com\/|spotify:)artist(\/|:)\w*"
        )

    async def processQuery(
        self, interaction: Interaction, query: str
    ) -> list[Entry] | None:
        if (
            self.youtube_video_regex.match(query) is not None
            or self.youtube_video_short_regex.match(query) is not None
        ):
            data = await self.youtube.fetchData(query)
            if data is None:
                return None

            with open("search_result.json", "w") as file:
                file.write(json.dumps(data))

            entry = Entry(
                title=Title(name=data["title"], url=data["webpage_url"]),
                artist=Artist(name=data["channel"], url=data["channel_url"]),
                duration=data["duration"],
                thumbnail=data["thumbnail"],
                requesterId=interaction.user.id,
                source=data["webpage_url"],
            )

            await self.fileManager.download(interaction, [entry])

            return [entry]

        if self.youtube_playlist_regex.match(query) is not None:
            data = await self.youtube.getData(query)
            if data is None:
                return None

            with open("search_result.json", "w") as file:
                file.write(json.dumps(data))

            playlist = Playlist(
                title=data["title"],
                url=data["webpage_url"],
                owner=Artist(name=data["channel"], url=data["channel_url"]),
            )

            return [
                Entry(
                    title=Title(name=data["title"], url=data["webpage_url"]),
                    artist=Artist(name=data["channel"], url=data["channel_url"]),
                    duration=data["duration"],
                    thumbnail=data["thumbnail"],
                    requesterId=interaction.user.id,
                    playlist=playlist,
                    source=data["url"],
                )
                for data in data["entries"]
            ]

        if self.spotify_track_regex.match(query) is not None:
            data = await self.spotify.getTrack(query)

            with open("spotify_track_result.json", "w") as file:
                file.write(json.dumps(data))

            return [
                Entry(
                    title=Title(
                        name=data["name"], url=data["external_urls"]["spotify"]
                    ),
                    artist=Artist(
                        name=data["artists"][0]["name"],
                        url=data["artists"][0]["external_urls"]["spotify"],
                    ),
                    album=Album(
                        name=data["album"]["name"],
                        url=data["album"]["external_urls"]["spotify"],
                    ),
                    duration=0,
                    thumbnail=data["album"]["images"][0]["url"],
                    requesterId=interaction.user.id,
                )
            ]

        if (
            self.spotify_playlist_regex.match(query) is not None
            or self.spotify_album_regex.match(query) is not None
            or self.spotify_artist_regex.match(query) is not None
        ):
            return

        return await self.processQuery(interaction, await self.search(query))

    async def search(self, query) -> str:
        result = await self.youtube.searchVideo(query)
        return result["result"][0]["link"]