From c68edc34c952c76a78c575aa2faa260e4a08d2f5 Mon Sep 17 00:00:00 2001
From: Lightemerald <light_emerald@aostia.com>
Date: Wed, 15 Nov 2023 11:07:16 +0100
Subject: [PATCH] Fixes

---
 setupScript.ps1 | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/setupScript.ps1 b/setupScript.ps1
index 926e490..836d828 100644
--- a/setupScript.ps1
+++ b/setupScript.ps1
@@ -51,7 +51,14 @@ function Setup {
     ChocoInstallApps -Type $Type
     InstallApps -Type $Type
     UpdateWindows
+    Show-WDACMenu
     Write-Host "Installation complete!"
+    Write-Host "Would you like to restart now? (y/n)"
+    $input = Read-Host
+    if ($input -eq "y") {
+        Clear-RecycleBin
+        Restart-Computer -Force
+    }
 }
 
 function Custom {
@@ -138,6 +145,13 @@ function SetupCustom {
     if ($Update) {
         UpdateWindows
     }
+    Show-WDACMenu
+    Write-Host "Would you like to restart now? (y/n)"
+    $input = Read-Host
+    if ($input -eq "y") {
+        Clear-RecycleBin
+        Restart-Computer -Force
+    }
 }
 
 function Show-InstallationMenu {
@@ -167,12 +181,4 @@ function Show-InstallationMenu {
     } while ($selection -ne "1" -and $selection -ne "2" -and $selection -ne "3")
 }
 
-Show-InstallationMenu
-Show-WDACMenu
-
-Write-Host "Would you like to restart now? (y/n)"
-$input = Read-Host
-if ($input -eq "y") {
-    Clear-RecycleBin
-    Restart-Computer -Force
-}
\ No newline at end of file
+Show-InstallationMenu
\ No newline at end of file