From 18219274713a81683d91bcbb7d7c0bffd3ca03af Mon Sep 17 00:00:00 2001
From: Abhi <allinoneallinone00@gmail.com>
Date: Tue, 6 Aug 2024 16:41:55 +0530
Subject: [PATCH] fix ytdl, unload musicbot

---
 modules/loader.py |  6 ++++++
 utils/scripts.py  | 25 ++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/modules/loader.py b/modules/loader.py
index b6a80fe..868aca1 100644
--- a/modules/loader.py
+++ b/modules/loader.py
@@ -7,6 +7,9 @@
 #  (at your option) any later version.
 import hashlib
 import os
+import shutil
+import subprocess
+import sys
 
 import requests
 from pyrogram import Client, filters
@@ -141,6 +144,9 @@ async def unload_mods(_, message: Message):
 
     if os.path.exists(f"{BASE_PATH}/modules/custom_modules/{module_name}.py"):
         os.remove(f"{BASE_PATH}/modules/custom_modules/{module_name}.py")
+        if module_name == "musicbot":
+            subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "requirements.txt"], cwd=f"{BASE_PATH}/musicbot")
+            shutil.rmtree(f"{BASE_PATH}/musicbot")
         await message.edit(f"<b>The module <code>{module_name}</code> removed!</b>")
         restart()
     elif os.path.exists(f"{BASE_PATH}/modules/{module_name}.py"):
diff --git a/utils/scripts.py b/utils/scripts.py
index 94f7922..8ff551e 100644
--- a/utils/scripts.py
+++ b/utils/scripts.py
@@ -29,10 +29,13 @@
 from typing import Dict, Tuple
 
 from PIL import Image
+import psutil
 from pyrogram import Client, errors, enums
 from pyrogram.errors import FloodWait, MessageNotModified
 from pyrogram.types import Message
 
+from utils import db
+
 from .misc import modules_help, prefix, requirements_list
 
 META_COMMENTS = re.compile(r"^ *# *meta +(\S+) *: *(.*?)\s*$", re.MULTILINE)
@@ -211,6 +214,13 @@ def text(message: Message) -> str:
 
 
 def restart() -> None:
+    music_bot_pid = db.get("custom.musicbot", "music_bot_pid", None)
+    if music_bot_pid is not None:
+        try:
+            music_bot_process = psutil.Process(music_bot_pid)
+            music_bot_process.terminate()
+        except psutil.NoSuchProcess:
+            print("Music bot is not running.")
     os.execvp(sys.executable, [sys.executable, "main.py"])
 
 
@@ -318,7 +328,7 @@ def import_library(library_name: str, package_name: str = None):
         return importlib.import_module(library_name)
     except ImportError as exc:
         completed = subprocess.run(
-            [sys.executable, "-m", "pip", "install", package_name], check=True)
+            [sys.executable, "-m", "pip", "install", "--upgrade", package_name], check=True)
         if completed.returncode != 0:
             raise AssertionError(
                 f"Failed to install library {package_name} (pip exited with code {completed.returncode})"
@@ -326,6 +336,19 @@ def import_library(library_name: str, package_name: str = None):
         return importlib.import_module(library_name)
 
 
+def uninstall_library(package_name: str):
+    """
+    Uninstalls a library
+    :param package_name: package name in PyPi (pip uninstall example)
+    """
+    completed = subprocess.run(
+        [sys.executable, "-m", "pip", "uninstall", "-y", package_name], check=True)
+    if completed.returncode != 0:
+        raise AssertionError(
+            f"Failed to uninstall library {package_name} (pip exited with code {completed.returncode})"
+        )
+
+
 def resize_image(
     input_img, output=None, img_type="PNG", size: int = 512, size2: int = None
 ):