diff --git a/source/main.c b/source/main.c
index eb378b1..0a0243b 100644
--- a/source/main.c
+++ b/source/main.c
@@ -3,9 +3,11 @@
 #include <stdlib.h>
 #include <switch.h>
 #include "save_basic.h"
+#include "ui_basic.h"
 
 int main(int argc, char **argv)
 {
+    setvbuf(stdout, NULL, _IONBF, BUFSIZ);
     consoleInit(NULL);
     
     // Configure input
@@ -13,17 +15,22 @@ int main(int argc, char **argv)
     PadState pad;
     padInitializeDefault(&pad);
     
-    printf("BOTW Save Tool v0.2\n");
-    printf("Press A to test save loading\n");
-    printf("Press + to exit\n");
+    // Initialize UI
+    initUI();
+    
+    printMessage("Press A to test save loading\n");
+    printMessage("Press + to exit\n");
+    
+    bool exit = false;
     
     // Main loop
-    while(appletMainLoop())
+    while(appletMainLoop() && !exit)
     {
         padUpdate(&pad);
         u64 kDown = padGetButtonsDown(&pad);
         
-        if (kDown & HidNpadButton_Plus) break;
+        exit = handleButtonPress(kDown);
+        
         if (kDown & HidNpadButton_A) {
             openSave();
             processSave();