magi-code 0.63.0

use anyhow::Result;
use crossterm::{
    event::{
        DisableBracketedPaste, DisableMouseCapture, EnableBracketedPaste, EnableMouseCapture,
        KeyboardEnhancementFlags, PopKeyboardEnhancementFlags, PushKeyboardEnhancementFlags,
    },
    execute,
    terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode},
};
use std::io;

const ENHANCED_KEYBOARD_FLAGS: KeyboardEnhancementFlags =
    KeyboardEnhancementFlags::DISAMBIGUATE_ESCAPE_CODES
        .union(KeyboardEnhancementFlags::REPORT_EVENT_TYPES)
        .union(KeyboardEnhancementFlags::REPORT_ALL_KEYS_AS_ESCAPE_CODES);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TerminalOp {
    EnableRawMode,
    EnterAlternateScreenEnableInput,
    EnableEnhancedKeyboard,
    DisableEnhancedKeyboard,
    DisableInputAndLeaveAlternateScreen,
    DisableRawMode,
}

trait TerminalOps {
    fn run(&self, op: TerminalOp) -> Result<()>;

    fn enhanced_keyboard_supported(&self) -> bool {
        true
    }
}

#[derive(Debug, Default)]
struct CrosstermTerminalOps;

pub(crate) struct TerminalGuard {
    ops: Box<dyn TerminalOps>,
    enhanced_keyboard_enabled: bool,
}

impl TerminalOps for CrosstermTerminalOps {
    fn run(&self, op: TerminalOp) -> Result<()> {
        match op {
            TerminalOp::EnableRawMode => enable_raw_mode().map_err(Into::into),
            TerminalOp::EnterAlternateScreenEnableInput => execute!(
                io::stdout(),
                EnterAlternateScreen,
                EnableMouseCapture,
                EnableBracketedPaste
            )
            .map_err(Into::into),
            TerminalOp::EnableEnhancedKeyboard => execute!(
                io::stdout(),
                PushKeyboardEnhancementFlags(ENHANCED_KEYBOARD_FLAGS)
            )
            .map_err(Into::into),
            TerminalOp::DisableEnhancedKeyboard => {
                execute!(io::stdout(), PopKeyboardEnhancementFlags).map_err(Into::into)
            }
            TerminalOp::DisableInputAndLeaveAlternateScreen => execute!(
                io::stdout(),
                DisableBracketedPaste,
                DisableMouseCapture,
                LeaveAlternateScreen
            )
            .map_err(Into::into),
            TerminalOp::DisableRawMode => disable_raw_mode().map_err(Into::into),
        }
    }

    fn enhanced_keyboard_supported(&self) -> bool {
        !cfg!(windows)
    }
}

impl TerminalGuard {
    pub(crate) fn enter() -> Result<Self> {
        Self::enter_with_ops(Box::new(CrosstermTerminalOps))
    }

    fn enter_with_ops(ops: Box<dyn TerminalOps>) -> Result<Self> {
        ops.run(TerminalOp::EnableRawMode)?;
        if let Err(error) = ops.run(TerminalOp::EnterAlternateScreenEnableInput) {
            let _ = ops.run(TerminalOp::DisableInputAndLeaveAlternateScreen);
            let _ = ops.run(TerminalOp::DisableRawMode);
            return Err(error);
        }
        let enhanced_keyboard_supported = ops.enhanced_keyboard_supported();
        if enhanced_keyboard_supported
            && let Err(error) = ops.run(TerminalOp::EnableEnhancedKeyboard)
        {
            let _ = ops.run(TerminalOp::DisableEnhancedKeyboard);
            let _ = ops.run(TerminalOp::DisableInputAndLeaveAlternateScreen);
            let _ = ops.run(TerminalOp::DisableRawMode);
            return Err(error);
        }
        Ok(Self {
            ops,
            enhanced_keyboard_enabled: enhanced_keyboard_supported,
        })
    }
}

pub(crate) fn restore_terminal_best_effort() {
    let ops = CrosstermTerminalOps;
    let _ = ops.run(TerminalOp::DisableEnhancedKeyboard);
    let _ = ops.run(TerminalOp::DisableInputAndLeaveAlternateScreen);
    let _ = ops.run(TerminalOp::DisableRawMode);
}

impl Drop for TerminalGuard {
    fn drop(&mut self) {
        if self.enhanced_keyboard_enabled {
            let _ = self.ops.run(TerminalOp::DisableEnhancedKeyboard);
        }
        let _ = self
            .ops
            .run(TerminalOp::DisableInputAndLeaveAlternateScreen);
        let _ = self.ops.run(TerminalOp::DisableRawMode);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Debug)]
    struct FakeTerminalOps {
        calls: std::cell::RefCell<Vec<TerminalOp>>,
        fail_on: Option<TerminalOp>,
        enhanced_keyboard_supported: bool,
    }

    impl Default for FakeTerminalOps {
        fn default() -> Self {
            Self {
                calls: std::cell::RefCell::new(Vec::new()),
                fail_on: None,
                enhanced_keyboard_supported: true,
            }
        }
    }

    impl FakeTerminalOps {
        fn with_failure(fail_on: TerminalOp) -> Self {
            Self {
                fail_on: Some(fail_on),
                ..Self::default()
            }
        }

        fn without_enhanced_keyboard_support() -> Self {
            Self {
                enhanced_keyboard_supported: false,
                ..Self::default()
            }
        }
    }

    impl TerminalOps for std::rc::Rc<FakeTerminalOps> {
        fn run(&self, op: TerminalOp) -> Result<()> {
            self.calls.borrow_mut().push(op);
            if self.fail_on == Some(op) {
                anyhow::bail!("injected terminal failure: {op:?}");
            }
            Ok(())
        }

        fn enhanced_keyboard_supported(&self) -> bool {
            self.enhanced_keyboard_supported
        }
    }

    #[test]
    fn enhanced_keyboard_flags_request_event_types_for_standalone_ctrl() {
        assert!(
            ENHANCED_KEYBOARD_FLAGS.contains(KeyboardEnhancementFlags::DISAMBIGUATE_ESCAPE_CODES)
        );
        assert!(ENHANCED_KEYBOARD_FLAGS.contains(KeyboardEnhancementFlags::REPORT_EVENT_TYPES));
        assert!(
            ENHANCED_KEYBOARD_FLAGS
                .contains(KeyboardEnhancementFlags::REPORT_ALL_KEYS_AS_ESCAPE_CODES)
        );
    }

    #[test]
    fn successful_enter_drop_restores_terminal_state() {
        let ops = std::rc::Rc::new(FakeTerminalOps::default());
        {
            let _guard = TerminalGuard::enter_with_ops(Box::new(ops.clone())).unwrap();
        }

        assert_eq!(
            *ops.calls.borrow(),
            vec![
                TerminalOp::EnableRawMode,
                TerminalOp::EnterAlternateScreenEnableInput,
                TerminalOp::EnableEnhancedKeyboard,
                TerminalOp::DisableEnhancedKeyboard,
                TerminalOp::DisableInputAndLeaveAlternateScreen,
                TerminalOp::DisableRawMode,
            ]
        );
    }

    #[test]
    fn unsupported_enhanced_keyboard_starts_without_push_or_pop() {
        let ops = std::rc::Rc::new(FakeTerminalOps::without_enhanced_keyboard_support());
        {
            let _guard = TerminalGuard::enter_with_ops(Box::new(ops.clone())).unwrap();
        }

        assert_eq!(
            *ops.calls.borrow(),
            vec![
                TerminalOp::EnableRawMode,
                TerminalOp::EnterAlternateScreenEnableInput,
                TerminalOp::DisableInputAndLeaveAlternateScreen,
                TerminalOp::DisableRawMode,
            ]
        );
    }

    #[test]
    fn failed_enhanced_keyboard_setup_restores_terminal_state() {
        let ops = std::rc::Rc::new(FakeTerminalOps::with_failure(
            TerminalOp::EnableEnhancedKeyboard,
        ));

        let error = TerminalGuard::enter_with_ops(Box::new(ops.clone()))
            .err()
            .unwrap();

        assert!(error.to_string().contains("injected terminal failure"));
        assert_eq!(
            *ops.calls.borrow(),
            vec![
                TerminalOp::EnableRawMode,
                TerminalOp::EnterAlternateScreenEnableInput,
                TerminalOp::EnableEnhancedKeyboard,
                TerminalOp::DisableEnhancedKeyboard,
                TerminalOp::DisableInputAndLeaveAlternateScreen,
                TerminalOp::DisableRawMode,
            ]
        );
    }

    #[test]
    fn failed_alternate_screen_setup_cleans_up_raw_mode_before_returning_error() {
        let ops = std::rc::Rc::new(FakeTerminalOps::with_failure(
            TerminalOp::EnterAlternateScreenEnableInput,
        ));

        let error = TerminalGuard::enter_with_ops(Box::new(ops.clone()))
            .err()
            .unwrap();

        assert!(error.to_string().contains("injected terminal failure"));
        assert_eq!(
            *ops.calls.borrow(),
            vec![
                TerminalOp::EnableRawMode,
                TerminalOp::EnterAlternateScreenEnableInput,
                TerminalOp::DisableInputAndLeaveAlternateScreen,
                TerminalOp::DisableRawMode,
            ]
        );
    }
}