From f1b31d31435f8d50545441e59929dd40aa561d55 Mon Sep 17 00:00:00 2001
From: Choco <94597336+ChocoMeow@users.noreply.github.com>
Date: Sun, 26 May 2024 15:57:25 +0800
Subject: [PATCH] Rewritten some method in ipc_client

---
 ipc/client.py  |   1 -
 ipc/methods.py | 106 +++++++++++++++++++++++++++++++++----------------
 2 files changed, 71 insertions(+), 36 deletions(-)

diff --git a/ipc/client.py b/ipc/client.py
index 862d178..946657f 100644
--- a/ipc/client.py
+++ b/ipc/client.py
@@ -1,6 +1,5 @@
 import aiohttp
 import asyncio
-import traceback
 import logging
 
 from discord.ext import commands
diff --git a/ipc/methods.py b/ipc/methods.py
index f07f79c..0c4b129 100644
--- a/ipc/methods.py
+++ b/ipc/methods.py
@@ -2,10 +2,11 @@ import function as func
 
 from typing import (
     List,
-    Dict
+    Dict,
+    Union
 )
 
-from discord import Member, VoiceChannel
+from discord import User, Member, VoiceChannel
 from discord.ext import commands
 from voicelink import Player, Track, Playlist, NodePool, decode, LoopType
 
@@ -58,22 +59,29 @@ async def connect_channel(member: Member, bot: commands.Bot):
     except:
         return
 
-async def initUser(bot: commands.Bot, data: Dict) -> Dict:
+async def initBot(bot: commands.Bot, data: Dict) -> Dict:
     user_id = data.get("user_id")
     user = bot.get_user(user_id)
-    if not user: 
-        return
+    if not user:
+        user = await bot.fetch_user(user_id)
+
+    if user:
+        return {
+            "op": "initBot",
+            "user_id": user_id,
+            "bot_name": bot.user.display_name,
+            "bot_avatar": bot.user.display_avatar.url,
+            "bot_id": bot.user.id
+        }
+    
+async def initUser(bot: commands.Bot, data: Dict) -> Dict:
+    user_id = data.get("user_id")
+    data = await func.get_user(user_id)
     
-    guild_id = None
-    for guild in user.mutual_guilds:
-        member = guild.get_member(user.id)
-        if member.voice and member.voice.channel:
-            guild_id = guild.id
-            
     return {
         "op": "initUser",
         "user_id": user_id,
-        "guild_id": guild_id
+        "data": data
     }
     
 async def initPlayer(player: Player, member: Member, data: Dict):
@@ -98,6 +106,22 @@ async def initPlayer(player: Player, member: Member, data: Dict):
         "autoplay": player.settings.get("autoplay", False)
     }
 
+async def getRecommendation(bot: commands.Bot, data: Dict): 
+    node = NodePool.get_node()
+    if not node:
+        return
+    
+    track_data = decode(track_id := data.get("track_id"))
+    track = Track(track_id=track_id, info=track_data, requester=bot.user)
+    tracks: List[Track] = await node.get_recommendations(track)
+
+    return {
+        "op": "getRecommendation",
+        "user_id": data.get("user_id"),
+        "region": data.get("region"),
+        "tracks": [track.track_id for track in tracks] if tracks else []
+    }
+
 async def skipTo(player: Player, member: Member, data: Dict):
     if not player.is_privileged(member):
         if player.current and member == player.current.requester:
@@ -381,8 +405,10 @@ async def removePlaylistTrack(member: Member, data: Dict):
     
     await func.update_user(member.id, {"$pull": {f'playlist.{pId}.tracks': track_id }})
 
-methods: Dict[str, SystemMethod] = {
+methods: Dict[str, Union[SystemMethod, PlayerMethod, UserMethod]] = {
+    "initBot": SystemMethod(initBot),
     "initUser": SystemMethod(initUser),
+    "getRecommendation": SystemMethod(getRecommendation),
     "closeConnection": PlayerMethod(closeConnection),
     "initPlayer": PlayerMethod(initPlayer),
     "skipTo": PlayerMethod(skipTo),
@@ -409,32 +435,42 @@ async def process_methods(ipc_client, bot: commands.Bot, data: Dict) -> None:
         return
 
     try:
-        env: Dict = {}
+        env: Dict = {"bot": bot, "data": data}
         args: List = []
         
-        if guild_id := data.get("guild_id"):
-            if not (guild := bot.get_guild(guild_id)):
-                return
-            env["guild"] = guild
-                
         params = method.params
-        if "member" in params:
-            if not (guild := env.get("guild")) or not (member := guild.get_member(user_id)):
-                return
-            env["member"] = member
-            
-        if "player" in params:
-            if not (guild := env.get("guild")) or not (player := guild.voice_client):
-                if not method.auto_connect or not (member := env.get("member")):
+        if not (type(method) == SystemMethod):
+            if guild_id := data.get("guild_id"):
+                if (guild := bot.get_guild(guild_id)):
+                    env["guild"] = guild
+
+            else:
+                user: User = bot.get_user(user_id)
+                if not user:
                     return
-                player = await connect_channel(member, bot)
-            env["player"] = player
-        
-        if "bot" in params:
-            env["bot"] = bot
-        
-        if "data" in params:
-            env["data"] = data
+                
+                for guild in user.mutual_guilds:
+                    member = guild.get_member(user_id)
+                    if member.voice and member.voice.channel:
+                        env["guild"] = guild
+                        env["member"] = member
+                        break
+
+            if "member" in params and "member" not in env:
+                if not (guild := env.get("guild")) or not (member := guild.get_member(user_id)):
+                    return
+                env["member"] = member
+                
+            if "player" in params:
+                if not (guild := env.get("guild")) or not (player := guild.voice_client):
+                    if not method.auto_connect or not (member := env.get("member")):
+                        return
+                    player = await connect_channel(member, bot)
+
+                if player.channel.id != member.voice.channel.id:
+                    return
+                
+                env["player"] = player
         
         for param in params:
             args.append(env.get(param))