From cd647eb352eb1d3831da9f17ea180d3b823153dd Mon Sep 17 00:00:00 2001
From: Den Delimarsky <1389609+dend@users.noreply.github.com>
Date: Thu, 8 Apr 2021 16:30:56 -0700
Subject: [PATCH] Adding the option to have timed keep-awake

---
 .../espresso/Espresso.Shell/Core/APIHelper.cs | 37 ++++++++++++++++---
 .../espresso/Espresso.Shell/Program.cs        | 18 +++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/src/modules/espresso/Espresso.Shell/Core/APIHelper.cs b/src/modules/espresso/Espresso.Shell/Core/APIHelper.cs
index aa07896962..a8dd28a695 100644
--- a/src/modules/espresso/Espresso.Shell/Core/APIHelper.cs
+++ b/src/modules/espresso/Espresso.Shell/Core/APIHelper.cs
@@ -49,9 +49,9 @@ namespace Espresso.Shell.Core
             }
         }
 
-        public static bool SetIndefiniteKeepAwake(bool? keepDisplayOn = true)
+        public static bool SetIndefiniteKeepAwake(bool keepDisplayOn = true)
         {
-            if ((bool)keepDisplayOn)
+            if (keepDisplayOn)
             {
                 return SetAwakeState(EXECUTION_STATE.ES_SYSTEM_REQUIRED | EXECUTION_STATE.ES_DISPLAY_REQUIRED | EXECUTION_STATE.ES_CONTINUOUS);
             }
@@ -61,15 +61,42 @@ namespace Espresso.Shell.Core
             }    
         }
 
-        public static bool SetTimedKeepAwake(int seconds, bool keepDisplayOn = true)
+        public static bool SetTimedKeepAwake(long seconds, bool keepDisplayOn = true)
         {
             if (keepDisplayOn)
             {
-                return false;
+                var success = SetAwakeState(EXECUTION_STATE.ES_SYSTEM_REQUIRED | EXECUTION_STATE.ES_DISPLAY_REQUIRED | EXECUTION_STATE.ES_CONTINUOUS);
+                if (success)
+                {
+                    RunTimedLoop(seconds);
+                    return true;
+                }
+                else
+                {
+                    return false;
+                }
             }
             else
             {
-                return false;
+                var success = SetAwakeState(EXECUTION_STATE.ES_SYSTEM_REQUIRED | EXECUTION_STATE.ES_CONTINUOUS);
+                if (success)
+                {
+                    RunTimedLoop(seconds);
+                    return true;
+                }
+                else
+                {
+                    return false;
+                }
+            }
+        }
+
+        private static void RunTimedLoop(long seconds)
+        {
+            var startTime = DateTime.UtcNow;
+            while (DateTime.UtcNow - startTime < TimeSpan.FromSeconds(seconds))
+            {
+                // We do nothing.
             }
         }
     }
diff --git a/src/modules/espresso/Espresso.Shell/Program.cs b/src/modules/espresso/Espresso.Shell/Program.cs
index db8106c524..75fbe48d46 100644
--- a/src/modules/espresso/Espresso.Shell/Program.cs
+++ b/src/modules/espresso/Espresso.Shell/Program.cs
@@ -82,6 +82,24 @@ namespace Espresso.Shell
                     Console.WriteLine("Could not set up the state to be indefinite keep awake.");
                 }
             }
+            else
+            {
+                // Timed keep-awake.
+                bool success = APIHelper.SetTimedKeepAwake(timeLimit, displayOn);
+                if (success)
+                {
+                    Console.WriteLine($"Finished execution of timed keep-awake.");
+
+                    // Because the timed keep-awake execution completed, there is no reason for
+                    // Espresso to stay alive - I will just shut down the application until it's
+                    // launched again by the user.
+                    Environment.Exit(0);
+                }
+                else
+                {
+                    Console.WriteLine("Could not set up the state to be timed keep awake.");
+                }
+            }
 
             new ManualResetEvent(false).WaitOne();
         }