diff --git a/EDK2/build.sh b/EDK2/build.sh
index 92244db3..361d802e 100644
--- a/EDK2/build.sh
+++ b/EDK2/build.sh
@@ -1,23 +1,28 @@
 #!/bin/sh
 
-if [ -z "$1" ]; then
-    EDKARCH=X64
-    postfix=x64
-elif [ "$1" = "ia32" ]; then
-    EDKARCH=IA32
-    postfix=ia32
-    shift
-elif [ "$1" = "aa64" ]; then
-    EDKARCH=AARCH64
-    postfix=aa64
-    shift
-fi
+# Set architecture based on input argument
+case "$1" in
+    "ia32")
+        EDKARCH=IA32
+        postfix=ia32
+        ;;
+    "aa64")
+        EDKARCH=AARCH64
+        postfix=aa64
+        ;;
+    *)
+        EDKARCH=X64
+        postfix=x64
+        ;;
+esac
 
-cd edk2-edk2-stable201911
+cd edk2-edk2-stable201911 || exit 1
 
+# Clean up configuration and cache files
 rm -rf ./Conf/.cache
 rm -f ./Conf/.AutoGenIdFile.txt
 
+# Define paths for EFI files
 VTEFI_PATH=Build/MdeModule/RELEASE_GCC48/$EDKARCH/MdeModulePkg/Application/Ventoy/Ventoy/OUTPUT/Ventoy.efi
 DST_PATH=../../INSTALL/ventoy/ventoy_${postfix}.efi
 
@@ -27,17 +32,15 @@ DST_PATH2=../../INSTALL/ventoy/vtoyutil_${postfix}.efi
 VTEFI_PATH3=Build/MdeModule/RELEASE_GCC48/$EDKARCH/MdeModulePkg/Application/VDiskChain/VDiskChain/OUTPUT/VDiskChain.efi
 DST_PATH3=../../VDiskChain/Tool/vdiskchain_${postfix}.efi
 
-
-rm -f $VTEFI_PATH
-rm -f $DST_PATH
-rm -f $VTEFI_PATH2
-rm -f $DST_PATH2
-rm -f $VTEFI_PATH3
+# Remove old EFI files
+rm -f $VTEFI_PATH $DST_PATH $VTEFI_PATH2 $DST_PATH2 $VTEFI_PATH3
 [ -d ../../VDiskChain ] && rm -f $DST_PATH3
 
+# Setup build environment
 unset WORKSPACE
 source ./edksetup.sh
 
+# Build based on architecture
 if [ "$EDKARCH" = "AARCH64" ]; then    
     GCC48_AARCH64_PREFIX=aarch64-linux-gnu- \
     build -p MdeModulePkg/MdeModulePkg.dsc -a $EDKARCH -b RELEASE -t GCC48
@@ -45,6 +48,7 @@ else
     build -p MdeModulePkg/MdeModulePkg.dsc -a $EDKARCH -b RELEASE -t GCC48
 fi
 
+# Check if build was successful and copy files
 if [ -e $VTEFI_PATH ] && [ -e $VTEFI_PATH2 ] && [ -e $VTEFI_PATH3 ]; then
     echo -e '\n\n====================== SUCCESS ========================\n\n'    
     cp -a $VTEFI_PATH $DST_PATH
@@ -56,4 +60,3 @@ else
     cd ..
     exit 1
 fi
-