use std::io::{self, Write};
use crate::utils::tty::TtyExt;
use anyhow::{Context, Result};
use ratatui::crossterm::{
cursor::MoveToColumn,
event::{DisableBracketedPaste, DisableFocusChange, EnableBracketedPaste, EnableFocusChange},
execute,
terminal::{
self, Clear, ClearType, EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode,
enable_raw_mode,
},
};
#[derive(Debug)]
struct TerminalState {
raw_mode_enabled: bool,
bracketed_paste_enabled: bool,
focus_change_enabled: bool,
}
pub struct AlternateScreenSession {
original_state: TerminalState,
entered: bool,
}
impl AlternateScreenSession {
pub fn enter() -> Result<Self> {
let mut stdout = io::stdout();
let is_tty = stdout.is_tty_ext();
if !is_tty {
tracing::warn!("stdout is not a TTY, alternate screen features may not work");
}
let original_state = TerminalState {
raw_mode_enabled: false, bracketed_paste_enabled: false,
focus_change_enabled: false,
};
execute!(stdout, EnterAlternateScreen)
.context("failed to enter alternate screen for terminal app")?;
let mut session = Self {
original_state,
entered: true,
};
enable_raw_mode().context("failed to enable raw mode for terminal app")?;
session.original_state.raw_mode_enabled = true;
if is_tty && execute!(stdout, EnableBracketedPaste).is_ok() {
session.original_state.bracketed_paste_enabled = true;
}
if is_tty && execute!(stdout, EnableFocusChange).is_ok() {
session.original_state.focus_change_enabled = true;
}
Ok(session)
}
pub fn exit(mut self) -> Result<()> {
self.restore_state()?;
self.entered = false; Ok(())
}
pub fn run<F, T>(f: F) -> Result<T>
where
F: FnOnce() -> Result<T>,
{
let session = Self::enter()?;
let result = f();
session.exit()?;
result
}
fn restore_state(&mut self) -> Result<()> {
if !self.entered {
return Ok(());
}
while let Ok(true) = crossterm::event::poll(std::time::Duration::from_millis(0)) {
let _ = crossterm::event::read();
}
let mut stdout = io::stdout();
let _ = execute!(stdout, MoveToColumn(0), Clear(ClearType::CurrentLine));
let mut errors = Vec::new();
if let Err(e) = execute!(stdout, LeaveAlternateScreen) {
tracing::warn!(%e, "failed to leave alternate screen");
errors.push(format!("leave alternate screen: {}", e));
}
if self.original_state.focus_change_enabled
&& let Err(e) = execute!(stdout, DisableFocusChange)
{
tracing::warn!(%e, "failed to disable focus change");
errors.push(format!("disable focus change: {}", e));
}
if self.original_state.bracketed_paste_enabled
&& let Err(e) = execute!(stdout, DisableBracketedPaste)
{
tracing::warn!(%e, "failed to disable bracketed paste");
errors.push(format!("disable bracketed paste: {}", e));
}
if self.original_state.raw_mode_enabled
&& let Err(e) = disable_raw_mode()
{
tracing::warn!(%e, "failed to disable raw mode");
errors.push(format!("disable raw mode: {}", e));
}
if let Err(e) = stdout.flush() {
tracing::warn!(%e, "failed to flush stdout");
errors.push(format!("flush stdout: {}", e));
}
if errors.is_empty() {
Ok(())
} else {
tracing::warn!(
errors = ?errors,
"some terminal operations failed during restore"
);
Ok(())
}
}
}
impl Drop for AlternateScreenSession {
fn drop(&mut self) {
if self.entered {
let _ = self.restore_state();
}
}
}
pub fn clear_screen() -> Result<()> {
execute!(io::stdout(), Clear(ClearType::All)).context("failed to clear alternate screen")
}
pub fn terminal_size() -> Result<(u16, u16)> {
terminal::size().context("failed to get terminal size")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_enter_exit_cycle() {
if !io::stdout().is_tty_ext() {
return;
}
let session = AlternateScreenSession::enter();
assert!(session.is_ok());
if let Ok(session) = session {
let result = session.exit();
assert!(result.is_ok());
}
}
#[test]
fn test_run_with_closure() {
if !io::stdout().is_tty_ext() {
return;
}
let result = AlternateScreenSession::run(|| {
Ok(42)
});
assert!(result.is_ok());
assert_eq!(result.unwrap(), 42);
}
#[test]
fn test_run_with_error() {
let result: Result<()> = AlternateScreenSession::run(|| Err(anyhow::anyhow!("test error")));
assert!(result.is_err());
}
#[test]
fn test_drop_cleanup() {
{
let _session = AlternateScreenSession::enter();
}
}
}