ShiteMusicBot/cogs/utils/mixplayer/mixqueue.py

# Discord Packages
from lavalink import AudioTrack

from collections import OrderedDict, deque
from itertools import chain, cycle, islice
from random import shuffle


def roundrobin(*iterables):
    """roundrobin('ABC', 'D', 'EF') --> A D E B F C"""
    # Recipe credited to George Sakkis
    num_active = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))


class MixQueue:
    def __init__(self):
        self.queues = OrderedDict()
        self.priority_queue = []
        self._history = deque(maxlen=11)  # 10 + current
        self.looping = False
        self.loop_offset = 0

    def __str__(self):
        tmp = ''
        for requester in self.queues.keys():
            user_queue = list(map(str, self.queues[requester]))
            tmp += f'{requester}: {user_queue}\n'
        return tmp

    def __bool__(self):
        return not self.empty

    def __iter__(self):
        global_queue = roundrobin(*[x for x in self.queues.values()])
        priority_queue = self.priority_queue
        out = chain(priority_queue, global_queue)
        return out

    def __len__(self):
        length = 0
        for _ in self:
            length += 1
        return length

    def get_queue(self):
        if self.looping:
            offset = self.loop_offset
            return list(self)[offset:] + list(self)[:offset]
        else:
            return list(self)

    def clear(self):
        self.queues = OrderedDict()
        self.priority_queue = []
        self.looping = False
        self.loop_offset = 0

    # if dual is true also returns global positions of tracks
    def get_user_queue(self, requester: int, dual: bool = False):
        queue = self.queues.get(requester, [])
        if dual and queue:
            pos = [self._loc_to_glob(requester, i) for i in range(len(queue))]
            if self.looping:
                pos = [(p + self.loop_offset) % len(self) for p in pos]
            combined = zip(queue, pos)
            return list(combined)
        return queue

    def pop_first(self):
        if self.priority_queue:
            next_track = self.priority_queue.pop(0)
            self._history.append(next_track)
            return next_track
        try:
            if self.looping:
                try:
                    queue = list(self)
                    next_track = queue[self.loop_offset]
                    self.loop_offset += 1
                    if self.loop_offset >= len(self):
                        self.loop_offset = 0
                    self._history.append(next_track)
                    return next_track
                except IndexError as e:
                    print(e)
                    print(self.loop_offset)
                    print(len(self))
                    self.loop_offset = min(self.loop_offset, len(self)-1)
            else:
                next_track = self.queues[self.first_queue].pop(0)
                self._shuffle()
                self._clear_empty()
                self._history.append(next_track)
                return next_track
        except KeyError:
            pass

    def add_track(self, requester: int, track: AudioTrack, pos: int = None):
        user_queue = self.queues.get(requester)
        if user_queue is None:
            self.queues[requester] = [track]
            localpos = 0
        elif pos is None:
            user_queue.append(track)
            localpos = len(user_queue) - 1
        else:
            user_queue.insert(pos, track)
            localpos = pos

        # Return info about track position
        return track, self._loc_to_glob(requester, localpos), localpos

    def add_next_track(self, track: AudioTrack):
        self.priority_queue.append(track)

    def remove_user_queue(self, requester: int):
        user_queue = self.queues.get(requester, [])
        if user_queue:
            if self.looping:
                # Should work, kinda hacky.
                for pos, track in reversed(list(enumerate(user_queue))):
                    self.remove_user_track(requester, pos)
            else:
                self.queues.pop(requester)

    def remove_user_track(self, requester: int, pos: int):
        user_queue = self.queues.get(requester)
        if user_queue is not None:
            if pos < len(user_queue):
                glob_pos = self._loc_to_glob(requester, pos)
                if glob_pos <= self.loop_offset:
                    self.loop_offset -= 1
                track = user_queue.pop(pos)
                self._clear_empty()
                return track

    def remove_global_track(self, pos: int):
        # Get the actual index of the song
        pos = list(self).index(self.get_queue()[pos])

        if pos <= self.loop_offset:
            self.loop_offset -= 1

        q, pos = self._glob_to_loc(pos)
        if q is None or pos is None:
            return
        queue = self.queues.get(q)
        track = queue.pop(pos)
        self._clear_empty()
        return track

    def move_user_track(self, requester: int, initial: int, final: int):
        queue = self.queues.get(requester, [])
        if queue:
            try:
                track = queue.pop(initial)
                queue.insert(final, track)
                return track
            except IndexError:
                pass

    def shuffle_user_queue(self, requester: int):
        queue = self.queues.get(requester, [])
        if queue:
            shuffle(queue)

    def enable_looping(self, looping):
        if (not self.looping) and looping:  # Not enabled and turning on
            self.looping = looping
            self.loop_offset = 0
        elif self.looping and not looping:  # Enabled to turning off
            tmp = [track for track in self.get_queue()]
            for key, _ in self.queues.items():
                self.queues = OrderedDict()
            self.looping = looping
            for track in tmp:
                self.add_track(track.requester, track)
            self.loop_offset = 0

    # Switches the order of user queues
    def _shuffle(self):
        self.queues.move_to_end(self.first_queue)

    # Removes empty user queues
    def _clear_empty(self):
        to_remove = [q for q in reversed(self.queues) if not self.queues[q]]
        for i in to_remove:
            self.queues.pop(i)

    def _loc_to_glob(self, requester, pos):
        globpos = len(self.priority_queue)
        passed = False
        for key, queue in self.queues.items():
            if len(queue) <= pos:
                globpos += len(queue)
            else:
                globpos += pos
                if key == requester:
                    passed = True
                if not passed:
                    globpos += 1
        return globpos

    def _glob_to_loc(self, pos: int):
        try:
            song = next(islice(self, pos, pos + 1))
        except (ValueError, StopIteration):
            return None, None
        user_queue = self.queues.get(song.requester, [])

        # In case song is in the priority queue
        if not user_queue or song not in user_queue:
            return None, pos

        for i, s in enumerate(user_queue):
            if s == song:
                return song.requester, i

    @property
    def first_queue(self):
        try:
            return list(self.queues.keys())[0]
        except IndexError:
            pass

    @property
    def empty(self):
        return len(self) == 0

    @property
    def history(self):
        return list(reversed(self._history))