From b4c76106eb18cc1ea7e05d2542701ccc53b06b81 Mon Sep 17 00:00:00 2001
From: Hardel <42413417+Hardel-DW@users.noreply.github.com>
Date: Tue, 29 Jun 2021 01:24:43 +0200
Subject: [PATCH] Show/Remove ImportantTaskText in freeplay

---
 Harion/CustomRoles/FreeplayTaskTester/FileRole.cs | 6 +++---
 Harion/CustomRoles/RoleManager.cs                 | 7 ++++++-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/Harion/CustomRoles/FreeplayTaskTester/FileRole.cs b/Harion/CustomRoles/FreeplayTaskTester/FileRole.cs
index 60da341..201a7b1 100644
--- a/Harion/CustomRoles/FreeplayTaskTester/FileRole.cs
+++ b/Harion/CustomRoles/FreeplayTaskTester/FileRole.cs
@@ -67,12 +67,12 @@ private static void Start(TaskAddButton taskFolder, RoleManager role) {
         private static void OnClick(TaskAddButton taskFolder, RoleManager role) {
             GameData.PlayerInfo data = PlayerControl.LocalPlayer.Data;
             if (role.HasRole(data.Object)) {
+                role.RpcRemovePlayer(data.Object);
+                taskFolder.Overlay.enabled = false;
+
                 if (data.IsImpostor) {
                     PlayerControl.LocalPlayer.RemoveInfected();
                 }
-
-                role.RpcRemovePlayer(data.Object);
-                taskFolder.Overlay.enabled = false;
             } else {
                 if (RoleManager.HasMainRole(data.Object)) {
                     RoleManager Role = RoleManager.GetMainRole(data.Object);
diff --git a/Harion/CustomRoles/RoleManager.cs b/Harion/CustomRoles/RoleManager.cs
index 5039d3c..2a18ad4 100644
--- a/Harion/CustomRoles/RoleManager.cs
+++ b/Harion/CustomRoles/RoleManager.cs
@@ -135,8 +135,10 @@ private void AddPlayer(byte PlayerId) {
             if (Player != null) {
                 RoleManager MainRole = GetMainRole(Player);
 
-                if (MainRole == null || !IsMainRole)
+                if (MainRole == null || !IsMainRole) {
                     AllPlayers.Add(Player);
+                    AddImportantTasks(Player);
+                }
             }
 
             DefineVisibleByWhitelist();
@@ -185,6 +187,9 @@ public void RpcAddPlayerRange(List<byte> PlayerIds) {
 
         // Remove method
         private void RemovePlayer(byte PlayerId) {
+            if (PlayerId == PlayerControl.LocalPlayer.PlayerId)
+                RemoveImportantTasks(PlayerControl.LocalPlayer);
+
             AllPlayers.RemovePlayer(PlayerId);
             DefineVisibleByWhitelist();
         }