From 363268e61e2372d11b90cfbc8179435d3604036a Mon Sep 17 00:00:00 2001
From: Nolan <nolan@thewordnerd.info>
Date: Sun, 12 Feb 2023 10:44:26 -0600
Subject: [PATCH] fix: Restore original scheme on exit.

---
 Cargo.toml |  1 +
 src/lib.rs | 20 ++++++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index dd4a49d..b7224f6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,6 +10,7 @@ edition = "2021"
 
 [dependencies]
 bevy = { version = "0.9", default-features = false }
+ctrlc = { version = "3", features = ["termination"] }
 
 [target.'cfg(windows)'.dependencies]
 windows = { version = "0.44", features = ["Win32_Foundation", "Win32_System_Power", "Win32_System_Registry", "Win32_System_SystemServices"] }
diff --git a/src/lib.rs b/src/lib.rs
index 403fc4d..92593fe 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,4 @@
-use bevy::{prelude::*, window::WindowFocused};
+use bevy::{app::AppExit, prelude::*, window::WindowFocused};
 #[cfg(windows)]
 use windows::{
     core::GUID,
@@ -22,6 +22,11 @@ fn setup(mut commands: Commands) {
         if let Some(active) = active.as_ref() {
             let scheme = DefaultScheme(*active);
             commands.insert_resource(scheme);
+            ctrlc::set_handler(move || {
+                Power::PowerSetActiveScheme(None, Some(active));
+                std::process::exit(1);
+            })
+            .expect("Failed to set exit handler");
         }
     }
     #[cfg(not(windows))]
@@ -51,6 +56,16 @@ fn focus_change(
     }
 }
 
+#[allow(unused_variables)]
+fn exit(mut exit: EventReader<AppExit>, scheme: Res<DefaultScheme>) {
+    for event in exit.iter() {
+        #[cfg(windows)]
+        unsafe {
+            Power::PowerSetActiveScheme(None, Some(&**scheme));
+        }
+    }
+}
+
 #[derive(Resource, Clone, Copy, Default)]
 pub struct FullThrottlePlugin {
     pub restore_original_scheme_on_unfocus: bool,
@@ -60,6 +75,7 @@ impl Plugin for FullThrottlePlugin {
     fn build(&self, app: &mut App) {
         app.insert_resource(*self)
             .add_startup_system(setup)
-            .add_system(focus_change);
+            .add_system(focus_change)
+            .add_system_to_stage(CoreStage::PostUpdate, exit);
     }
 }