From 12b51bcf090e80803e2bab681bb91ea52e327df3 Mon Sep 17 00:00:00 2001
From: longpanda <admin@ventoy.net>
Date: Sat, 24 Dec 2022 09:39:27 +0800
Subject: [PATCH] Change the prompt format when check device failed. (#2094)

---
 .../grub-2.04/grub-core/ventoy/ventoy_cmd.c   | 77 +++++++++++++------
 1 file changed, 55 insertions(+), 22 deletions(-)

diff --git a/GRUB2/MOD_SRC/grub-2.04/grub-core/ventoy/ventoy_cmd.c b/GRUB2/MOD_SRC/grub-2.04/grub-core/ventoy/ventoy_cmd.c
index c4c75861..22ba948f 100644
--- a/GRUB2/MOD_SRC/grub-2.04/grub-core/ventoy/ventoy_cmd.c
+++ b/GRUB2/MOD_SRC/grub-2.04/grub-core/ventoy/ventoy_cmd.c
@@ -5005,40 +5005,73 @@ int ventoy_load_part_table(const char *diskname)
     return 0;
 }
 
-static grub_err_t ventoy_cmd_load_part_table(grub_extcmd_context_t ctxt, int argc, char **args)
+static void ventoy_prompt_end(void)
 {
-    int ret;
+    int op = 0;
     char c;
-    
-    (void)argc;
-    (void)ctxt;
 
-    ret = ventoy_load_part_table(args[0]);
-    if (ret)
-    {
-        grub_printf("\n\nPress r/h/e for the corresponding operation.\n");
-        grub_printf(" r --- Reboot\n");
-        grub_printf(" h --- Halt\n");
-        grub_printf(" e --- Exit grub\n");
-        grub_refresh();
+    grub_printf("\n\n\n");
+    grub_printf(" 1 --- Exit grub\n");
+    grub_printf(" 2 --- Reboot\n");
+    grub_printf(" 3 --- Shut down\n");
+    grub_printf("Please enter your choice: ");
+    grub_refresh();
 
-        while (1)
+    while (1)
+    {
+        c = grub_getkey();
+        if (c >= '1' && c <= '3')
         {
-            c = grub_getkey();
-            if (c == 'r')
+            if (op == 0)
             {
-                grub_reboot();
+                op = c - '0';
+                grub_printf("%c", c);
+                grub_refresh();
             }
-            else if (c == 'h')
-            {   
-                grub_script_execute_sourcecode("halt");
+        }
+        else if (c == '\r' || c == '\n')
+        {
+            if (op)
+            {
+                if (op == 1)
+                {
+                    grub_exit();
+                }
+                else if (op == 2)
+                {
+                    grub_reboot();
+                }
+                else if (op == 3)
+                {   
+                    grub_script_execute_sourcecode("halt");
+                }
             }
-            else if (c == 'e')
+        }
+        else if (c == '\b')
+        {
+            if (op)
             {
-                grub_exit();
+                op = 0;
+                grub_printf("\rPlease enter your choice:   ");
+                grub_printf("\rPlease enter your choice: ");
+                grub_refresh();
             }
         }
     }
+}
+
+static grub_err_t ventoy_cmd_load_part_table(grub_extcmd_context_t ctxt, int argc, char **args)
+{
+    int ret;
+    
+    (void)argc;
+    (void)ctxt;
+
+    ret = ventoy_load_part_table(args[0]);
+    if (ret)
+    {
+        ventoy_prompt_end();
+    }
 
     g_ventoy_disk_part_size[0] = ventoy_get_vtoy_partsize(0);
     g_ventoy_disk_part_size[1] = ventoy_get_vtoy_partsize(1);