cogs/tjadb/tjadb.py

import random
import requests

from redbot.core  import Config, commands, checks
from urllib.parse import urljoin

from .functions import search_songs, parse_difficulty, parse_comparison
from .embeds    import embed_search, embed_random, embed_sotd, embed_color

__PLUGIN_ID__  = 0x298f2a081ede4dc486d8cc65b848b599
__API_BROWSE__ = "/api/browse"
__API_SOTD__   = "/api/sotd"

class TJADB(commands.Cog):
    """TJADB Search cog"""

    def __init__(self, bot):
        self.bot = bot
        self.conf = Config.get_conf(self, identifier=__PLUGIN_ID__)
        self.conf.register_global(
            url          = "https://tjadataba.se",
            color_search = "4DF479",
            color_random = "8E9926",
            color_sotd   = "17D4A5"
        )
        self.conf.register_user(
            difficulty={}
        )
        self.colors = {'search': self.conf.color_search,
                       'random': self.conf.color_random,
                       'sotd':   self.conf.color_sotd}

        self.cache = []


    @commands.command()
    async def sotd(self, ctx):
        """Displays the "Song of the day" """
        try:
            data = requests.get(urljoin(await self.conf.url(), __API_SOTD__)).json()
        except Exception as e:
            print(e)
            await ctx.send("""I could not fetch todays "Song of the day" """)
            return
        embed = await embed_sotd(data, self)
        await ctx.send(embed = embed)


    @commands.command()
    async def search(self, ctx, data):
        """Searches TJADB for strings that match"""
        songs = search_songs(data, self.cache)
        embed = await embed_search(data, songs, self)
        await ctx.send(embed = embed)


    @commands.command()
    async def random(self, ctx, difficulty=None):
        """ Gives a random song from the database. Optional difficulty requirement"""
        difficulty = parse_difficulty(difficulty)
        songs      = self.cache
        if len(songs) == 0:
            await ctx.send("Sorry, I don't have any songs loaded yet.")
            return
        # Filter if difficulty is set
        if difficulty:
            songs = [s for s in songs if s[difficulty]]
        # Filter on preferences:
        pref = await self.conf.user(ctx.author).difficulty.all()
        # Search filter:
        def match_diff(song):
            for diff, range in pref.items():
                if song[diff] in range:
                    return True
            return False
        if pref:
            if difficulty:
                pref  = {difficulty: pref[difficulty]}
            songs = list(filter(match_diff, songs))
        # Check if list is not empty
        if len(songs) == 0:
            await ctx.send(f"I'm sorry, but no results match your query. You may want to expand your preferred difficulty range.")
            return
        song  = random.choice(songs)
        await ctx.send(embed = await embed_random(song, self) )


    @commands.group()
    async def preference(self, ctx):
        """ Base command for handling user preferences"""
        pass

    @preference.command(name="difficulty")
    async def preferences_difficulty(self, ctx, course, *data):
        """ Set the preferred difficulty (e.g "preference difficulty oni 2 < x < 7")"""
        diff = parse_difficulty(course)
        if not diff:
            await ctx.send(f"'{course}' is not a valid difficulty")
            return
        diff_p = diff[2:]
        # Show settings for difficulty
        if len(data) == 0:
            range = await self.conf.user(ctx.author).difficulty.get_attr(diff)()
            if range == None:
                await ctx.send(f"Preferred difficulty for {diff_p} not set.")
            else:
                await ctx.send(f"Preferred difficulty for {diff_p} is set to: {range[0]}-{range[-1]}")
            return
        data = " ".join(data)
        # clear data if requested
        if data in ['clear', 'reset', 'remove', 'delete']:
            await self.conf.user(ctx.author).difficulty.get_attr(diff).clear()
            await ctx.send(f"Preferred difficulty for {diff_p} has been cleared")
            return
        try:
            range = parse_comparison(data)
        except:
            await ctx.send(f"'{data}' is not a valid comparison string. \nAn example would be: '2 < x < 7'")
            return
        await self.conf.user(ctx.author).difficulty.get_attr(diff).set(range)
        await ctx.send(f"Preferred difficulty for {diff_p} set to: {range[0]}-{range[-1]}")


    @commands.group()
    @checks.admin_or_permissions(manage_guild=True)
    async def tjadb(self, ctx):
        """ Base command for handling admin commands"""
        pass


    @tjadb.command(name="update")
    async def tjadb_update(self, ctx, verbose=None):
        """ Updates song cache"""
        url = await self.conf.url()
        try:
            self.cache = requests.get(urljoin(url, __API_BROWSE__)).json()
        except Exception as e:
            print(e)
            await ctx.send("Could not fetch the API data. Please check URL")
            return
        if verbose and verbose.lower() in ["verbose", "v", "out", "output", "return", "echo"]:
            await ctx.send("Cache updated")


    @tjadb.command(name="color")
    async def tjadb_color(self, ctx, color, value):
        """ Set the color values for the embeds"""
        if color not in self.colors.keys():
            await ctx.send(f"Not a valid color. Valid colors: {', '.join(self.colors.keys())}")
            return
        if value.startswith("#"):
            value = value[1:]
        try:
            assert len(value) == 6
            int(value, 16)
        except:
            await ctx.send(f"{value} is not a valid color.")
        await self.colors[color].set(value)
        await ctx.send(f"Color set")


    @tjadb.command(name="colors")
    async def tjadb_colors(self, ctx):
        """ Shows the color values for the embeds"""
        for key, val in self.colors.items():
            color = await val()
            await ctx.send(embed=embed_color(key, color))


    @tjadb.command(name = "url")
    async def tjadb_url(self, ctx, data=None):
        if not data:
            url = await self.conf.url()
            await ctx.send(f"The current URL is: {url}")
            return
        await self.conf.url.set(data)
        await ctx.send(f"URL is set to: {data}")