2 files changed, 28 insertions, 4 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 51422a5..f455160 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "bevy_terminal_display"
-version = "0.4.1"
+version = "0.4.2"
 edition = "2021"
 license = "0BSD OR MIT OR Apache-2.0"
 description = "A plugin for the Bevy game engine which enables rendering to a terminal using unicode braille characters."
@@ -15,4 +15,5 @@ bevy = "0.14"
 bevy_headless_render = "0.1"
 bevy_dither_post_process = "0.2"
 ratatui = "0.26"
+color-eyre = "0.6"
 
diff --git a/src/lib.rs b/src/lib.rs
index bf08319..682fca2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,9 +3,7 @@
 //! Bevy plugin which allows a camera to render to a terminal window.
 
 use std::{
-    fs::OpenOptions,
-    path::PathBuf,
-    sync::{Arc, Mutex},
+    fs::OpenOptions, io::stdout, path::PathBuf, sync::{Arc, Mutex}
 };
 
 use bevy::{
@@ -17,7 +15,9 @@ use bevy::{
 use bevy_dither_post_process::DitherPostProcessPlugin;
 
 use bevy_headless_render::HeadlessRenderPlugin;
+use color_eyre::config::HookBuilder;
 pub use crossterm;
+use crossterm::{event::{DisableMouseCapture, PopKeyboardEnhancementFlags}, terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen}, ExecutableCommand};
 use once_cell::sync::Lazy;
 pub use ratatui;
 
@@ -67,6 +67,21 @@ impl Plugin for TerminalDisplayPlugin {
             .with_filter(EnvFilter::builder().parse_lossy(format!("{},{}", Level::INFO, "wgpu=error,naga=warn")));
         let subscriber = Registry::default().with(file_layer);
         subscriber::set_global_default(subscriber).unwrap();
+
+        let (panic, error) = HookBuilder::default().into_hooks();
+        let panic = panic.into_panic_hook();
+        let error = error.into_eyre_hook();
+
+        color_eyre::eyre::set_hook(Box::new(move |e| {
+            let _ = restore_terminal();
+            error(e)
+        })).unwrap();
+
+        std::panic::set_hook(Box::new(move |info| {
+            let _ = restore_terminal();
+            panic(info);
+        }));
+        
         app.add_plugins((
             DitherPostProcessPlugin,
             HeadlessRenderPlugin,
@@ -87,3 +102,11 @@ impl Plugin for TerminalDisplayPlugin {
         .add_event::<input::events::TerminalInputEvent>();
     }
 }
+
+fn restore_terminal() {
+    let mut stdout = stdout();
+    let _ = stdout.execute(PopKeyboardEnhancementFlags);
+    let _ = stdout.execute(DisableMouseCapture);
+    let _ = stdout.execute(LeaveAlternateScreen);
+    let _ = disable_raw_mode();
+}