1 files changed, 26 insertions, 3 deletions

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();
+}