1 files changed, 32 insertions, 19 deletions

diff --git a/src/ui/backend.rs b/src/ui/backend.rs
index b5b9833..4b5b124 100644
--- a/src/ui/backend.rs
+++ b/src/ui/backend.rs
@@ -6,7 +6,6 @@ use termion::raw::{IntoRawMode, RawTerminal};
 use termion::screen::AlternateScreen;
 use termion::screen::IntoAlternateScreen;
 
-#[cfg(feature = "mouse")]
 use termion::input::MouseTerminal;
 
 trait New {
@@ -15,24 +14,38 @@ trait New {
         Self: Sized;
 }
 
-#[cfg(feature = "mouse")]
-type Screen = MouseTerminal<AlternateScreen<RawTerminal<std::io::Stdout>>>;
-#[cfg(feature = "mouse")]
-impl New for Screen {
+pub enum Screen {
+    WithMouse(MouseTerminal<AlternateScreen<RawTerminal<std::io::Stdout>>>),
+    WithoutMouse(AlternateScreen<RawTerminal<std::io::Stdout>>),
+}
+
+impl Screen {
     // Returns alternate screen
-    fn new() -> io::Result<Self> {
+    fn new(mouse_support: bool) -> io::Result<Self> {
         let stdout = io::stdout().into_raw_mode()?;
-        Ok(MouseTerminal::from(stdout.into_alternate_screen().unwrap()))
+        if mouse_support {
+            Ok(Self::WithMouse(MouseTerminal::from(
+                stdout.into_alternate_screen().unwrap(),
+            )))
+        } else {
+            Ok(Self::WithoutMouse(stdout.into_alternate_screen().unwrap()))
+        }
     }
 }
-#[cfg(not(feature = "mouse"))]
-type Screen = AlternateScreen<RawTerminal<std::io::Stdout>>;
-#[cfg(not(feature = "mouse"))]
-impl New for Screen {
-    // Returns alternate screen
-    fn new() -> io::Result<Self> {
-        let stdout = std::io::stdout().into_raw_mode()?;
-        Ok(stdout.into_alternate_screen().unwrap())
+
+impl Write for Screen {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        match self {
+            Screen::WithMouse(t) => t.write(buf),
+            Screen::WithoutMouse(t) => t.write(buf),
+        }
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        match self {
+            Screen::WithMouse(t) => t.flush(),
+            Screen::WithoutMouse(t) => t.flush(),
+        }
     }
 }
 
@@ -44,8 +57,8 @@ pub struct AppBackend {
 }
 
 impl AppBackend {
-    pub fn new() -> io::Result<Self> {
-        let mut alt_screen = Screen::new()?;
+    pub fn new(mouse_support: bool) -> io::Result<Self> {
+        let mut alt_screen = Screen::new(mouse_support)?;
         // clears the screen of artifacts
         write!(alt_screen, "{}", termion::clear::All)?;
 
@@ -80,8 +93,8 @@ impl AppBackend {
         let _ = stdout().flush();
     }
 
-    pub fn terminal_restore(&mut self) -> io::Result<()> {
-        let mut new_backend = Self::new()?;
+    pub fn terminal_restore(&mut self, mouse_support: bool) -> io::Result<()> {
+        let mut new_backend = Self::new(mouse_support)?;
         std::mem::swap(&mut self.terminal, &mut new_backend.terminal);
         Ok(())
     }