shadow-terminal 0.2.1

//! Sending terminal output to end users. This "native" output is more useful when used directly
//! with Rust (as opposed to FFI users or STDOUT users).
//!
//! We try to be efficient in how we send output. If the change from the underlying PTY is small,
//! then we just send diffs. Otherwise we send the entire screen or scrollback buffer.

use std::io::Write as _;

use snafu::{OptionExt as _, ResultExt as _};
use termwiz::surface::Change as TermwizChange;
use termwiz::surface::Position as TermwizPosition;

/// Send raw output directly to the user's terminal without going via the renderer. Useful for
/// sending ANSI codes to the terminal emulator. For example, `^[?1h` to set "application mode"
/// arrow key codes.
///
/// # Errors
/// When writing to STDOUT fails.
#[inline]
pub fn raw_string_direct_to_terminal(
    string: &str,
) -> Result<(), crate::errors::ShadowTerminalError> {
    std::io::stdout()
        .write(string.as_bytes())
        .with_whatever_context(|err| {
            format!("Writing direct raw output to user's terminal: {err:?}")
        })?;
    std::io::stdout().flush().with_whatever_context(|err| {
        format!("Writing direct raw output to user's terminal: {err:?}")
    })
}

/// The mode of the terminal screen, therefore either the primary screen, where the scrollback is
/// collected, or the alternate screen, where apps like `vim`, `htop`, etc, get rendered.
#[derive(
    Clone, Debug, Default, serde::Serialize, serde::Deserialize, schemars::JsonSchema, Eq, PartialEq,
)]
#[non_exhaustive]
pub enum ScreenMode {
    /// The typical REPL mode of the terminal. Also can be thought of as a view onto the bottom of
    /// the scrollback.
    #[default]
    Primary,
    /// The so-called "alternate" screen where apps like `vim`, `htop`, etc, get displayed.
    Alternate,
}

/// Hopefully the most common form of output, therefore a small diff of changes.
#[derive(Clone)]
#[non_exhaustive]
pub enum SurfaceDiff {
    /// Output generated by the terminal whilst in REPL mode, aka, the "primary screen".
    Scrollback(ScrollbackDiff),
    /// The current view of the terminal, regardless of whether it's the primary or alternate
    /// screen.
    Screen(ScreenDiff),
}

/// The scrollback is a history, albeit limited, of all the output whilst in REPL mode, aka the
/// "primary screen".
///
/// Even though it's called the "scrollback", it's still the main interactive view, we just mostly
/// are seeing the bottom of the scrollback where the current prompt is.
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct ScrollbackDiff {
    /// A list of `termwiz` changes that, once applied, should bring any surfaces following this
    /// view, up to date.
    pub changes: Vec<TermwizChange>,
    /// The size of the underlying PTY at the time this diff was made.
    pub size: (usize, usize),
    /// The current position of the user's view on the scrollback. Is 0 when not scrolling.
    pub position: usize,
    /// The size of the current scrollback. Can increase up to the configured maximum.
    pub height: usize,
}

/// The constant view into the terminal, regardless of whether it's in primary or alternate screen.
///
/// However, sending diffs of it is only possible when in "alternate mode". This is because the
/// diffs generated by the scrollback only apply to the ever-growing scrollback buffer. The screen
/// on the other hand is always limited to a certain height, in which case diffs can't just be
/// additive, they must also replace what is under them.
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct ScreenDiff {
    /// A list of `termwiz` changes that, once applied, should bring any surfaces following this
    /// view, up to date.
    pub changes: Vec<TermwizChange>,
    /// Whether the terminal screen is primary or alternate.
    pub mode: ScreenMode,
    /// The size of the underlying PTY at the time this diff was made.
    pub size: (usize, usize),
    /// All the details about the user's cursor.
    pub cursor: wezterm_term::CursorPosition,
}

impl std::fmt::Debug for SurfaceDiff {
    #[expect(clippy::min_ident_chars, reason = "It's in the standard library")]
    #[inline]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let info = match self {
            Self::Scrollback(diff) => ("Scrollback", diff.changes.len(), diff.size),
            Self::Screen(diff) => ("Screen", diff.changes.len(), diff.size),
        };
        write!(
            f,
            "{} diff of {} change(s) {}x{}",
            info.0, info.1, info.2 .0, info.2 .1,
        )
    }
}

/// A complete, cell-for-cell duplicate of the current Wezterm shadow terminal.
///
/// When diffing is deemed ineffeicient, say when resizing, or when scrolling in `vim`, it's
/// hopefully more efficient to just send the entire view of the terminal. After all, diffs have to
/// be applied one by one, so there must come a point where it's cheaper to just send all the cell
/// data verbatim.
#[derive(Clone)]
#[non_exhaustive]
pub enum CompleteSurface {
    /// A complete scrollback.
    Scrollback(CompleteScrollback),
    /// The current view of the terminal, regardless of whether it's the primary or alternate
    /// screen.
    Screen(CompleteScreen),
}

impl std::fmt::Debug for CompleteSurface {
    #[expect(clippy::min_ident_chars, reason = "It's in the standard library")]
    #[inline]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let info = match self {
            Self::Scrollback(scrollback) => ("scrollback", &scrollback.surface),
            Self::Screen(screen) => ("screen", &screen.surface),
        };
        write!(
            f,
            "Complete {} surface: {}x{}",
            info.0,
            info.1.dimensions().0,
            info.1.dimensions().1
        )
    }
}

/// Every cell in the current scrollback, and the current location if the scrollback is actively
/// being scrolled.
#[derive(Default, Clone)]
#[non_exhaustive]
pub struct CompleteScrollback {
    /// The `termwiz` surface data.
    pub surface: termwiz::surface::Surface,
    /// The position of the current scroll, so it would be 0, if the user is not scrolling.
    pub position: usize,
}

/// Every cell in the current sreen, and the screen's mode.
#[derive(Default, Clone)]
#[non_exhaustive]
pub struct CompleteScreen {
    /// The `termwiz` surface data.
    pub surface: termwiz::surface::Surface,
    /// Whether the terminal is in primary or alternate mode.
    pub mode: ScreenMode,
}

impl CompleteScreen {
    /// Instantiate an empty `CompleteScreen`.
    #[inline]
    #[must_use]
    pub fn new(width: usize, height: usize) -> Self {
        Self {
            surface: termwiz::surface::Surface::new(width, height),
            mode: ScreenMode::default(),
        }
    }
}

#[derive(Clone, Debug)]
#[non_exhaustive]
/// All the possible kinds of output, whether they're primary, alternate, diffs or entire snapshots.
pub enum Output {
    /// A diff should be the most common output format, it's more efficient.
    Diff(SurfaceDiff),
    /// In certain cases, it's likely more efficient to just send all the cell data for the
    /// terminal. Or perhaps it's useful in moments of recovery or reset.
    Complete(CompleteSurface),
}

/// The kinds of surfaces that can be output.
#[derive(Debug)]
#[non_exhaustive]
pub enum SurfaceKind {
    /// The terminal scrollback, or "primary screen".
    Scrollback,
    /// The current view of the terminal, regardless of whether it's the primary or alternate
    /// screen.
    Screen,
}

impl Default for SurfaceDiff {
    #[inline]
    fn default() -> Self {
        Self::Scrollback(ScrollbackDiff::default())
    }
}

impl crate::shadow_terminal::ShadowTerminal {
    /// Build output for broadcasting to end users.
    pub(crate) fn build_current_output(
        &mut self,
        kind: &SurfaceKind,
    ) -> Result<Output, crate::errors::ShadowTerminalError> {
        tracing::trace!("Converting Wezterm terminal state to a `termwiz::surface::Surface`");

        let tty_size = self.terminal.get_size();
        let total_lines = self.terminal.screen().scrollback_rows();
        let changed_line_ids = self.terminal.screen().get_changed_stable_rows(
            0..total_lines.try_into().with_whatever_context(|err| {
                format!("Couldn't convert `total_lines` to `isize`: {err:?}")
            })?,
            self.last_sent.pty_sequence,
        );

        // TODO: Explore these heuristics. Maybe make them user configurable?
        let is_diff_efficient = match kind {
            SurfaceKind::Scrollback => changed_line_ids.len() < total_lines.div_euclid(2),
            SurfaceKind::Screen => changed_line_ids.len() < tty_size.rows,
        };

        let is_building_screen = matches!(kind, SurfaceKind::Screen);
        let is_resized = self.last_sent.pty_size != (tty_size.cols, tty_size.rows);
        let is_diff_possible = !is_resized && !is_building_screen;

        let output = if is_diff_efficient && is_diff_possible {
            self.build_diff(kind, changed_line_ids, tty_size, total_lines)?
        } else {
            self.build_complete_surface(kind, tty_size, total_lines)?
        };

        Ok(output)
    }

    /// Query the active terminal for its screen mode.
    fn get_screen_mode(&self) -> ScreenMode {
        if self.terminal.is_alt_screen_active() {
            ScreenMode::Alternate
        } else {
            ScreenMode::Primary
        }
    }

    /// Build a diff of the changes from the PTY
    fn build_diff(
        &mut self,
        kind: &SurfaceKind,
        changed_line_ids: Vec<wezterm_term::StableRowIndex>,
        tty_size: wezterm_term::TerminalSize,
        total_lines: usize,
    ) -> Result<Output, crate::errors::ShadowTerminalError> {
        tracing::trace!("Building diff from Wezterm for {kind:?} from lines: {changed_line_ids:?}");

        let changes = self.generate_changes(kind, Some(changed_line_ids))?;
        let diff = match kind {
            SurfaceKind::Scrollback => SurfaceDiff::Scrollback(ScrollbackDiff {
                changes,
                size: (tty_size.cols, tty_size.rows),
                position: self.scroll_position,
                height: total_lines,
            }),
            SurfaceKind::Screen => SurfaceDiff::Screen(ScreenDiff {
                mode: self.get_screen_mode(),
                changes,
                size: (tty_size.cols, tty_size.rows),
                cursor: self.terminal.cursor_pos(),
            }),
        };
        Ok(Output::Diff(diff))
    }

    /// Build an entire surface of all the cell data from the PTY.
    fn build_complete_surface(
        &mut self,
        kind: &SurfaceKind,
        tty_size: wezterm_term::TerminalSize,
        total_lines: usize,
    ) -> Result<Output, crate::errors::ShadowTerminalError> {
        tracing::trace!(
            "Building surface or diff from Wezterm for {kind:?} from lines: 0 to {total_lines:?}"
        );

        let changes = self.generate_changes(kind, None)?;
        let complete_surface = match kind {
            SurfaceKind::Scrollback => {
                let changes_count = changes.len();
                let mut surface = termwiz::surface::Surface::new(tty_size.cols, total_lines);
                surface.add_changes(changes);
                tracing::trace!(
                    "Sending complete Scrollback ({} changes): Sample:\n{:.100}\n...",
                    changes_count,
                    surface.screen_chars_to_string()
                );
                CompleteSurface::Scrollback(CompleteScrollback {
                    surface,
                    position: self.scroll_position,
                })
            }
            SurfaceKind::Screen => {
                let changes_count = changes.len();
                let mut surface = termwiz::surface::Surface::new(tty_size.cols, tty_size.rows);
                surface.add_changes(changes);
                tracing::trace!(
                    "Sending complete Screen ({}x{}, {} changes): Sample:\n{:.1000}\n...",
                    tty_size.cols,
                    tty_size.rows,
                    changes_count,
                    surface.screen_chars_to_string()
                );
                CompleteSurface::Screen(CompleteScreen {
                    surface,
                    mode: self.get_screen_mode(),
                })
            }
        };

        Ok(Output::Complete(complete_surface))
    }

    /// Generate a change set. It is used both for generating diffs and it is, perhaps
    /// surprisingly, the method required to construct an entire surface from scratch.
    fn generate_changes(
        &mut self,
        kind: &SurfaceKind,
        maybe_dirty_lines: Option<Vec<isize>>,
    ) -> Result<Vec<TermwizChange>, crate::errors::ShadowTerminalError> {
        let mut changes = Vec::new();
        let (line_ids, output_start) = self.calculate_line_ids(kind, maybe_dirty_lines)?;
        let screen = self.terminal.screen_mut();

        for line_id in line_ids {
            let line = screen.line_mut(line_id);
            let y = line_id - output_start;
            changes.push(TermwizChange::CursorPosition {
                x: TermwizPosition::Absolute(0),
                y: TermwizPosition::Absolute(y),
            });

            let mut wide_character_offset = 0;
            for cell in line.cells_mut() {
                // Wide characters, like say, "🤓", use up 2 cells in the terminal. The following
                // cell is always left blank. The Wezterm terminal already does this, and also adding
                // a wide character to a Termwiz surface will create these blank cells. Therefore
                // without intervention we'll actually create blank cells from both Wezterm and
                // Termwiz, doubling the number of needed blank cells. So we just ignore the blank
                // cells coming from Wezterm and let Termwiz handle automating all the blank cells.
                if wide_character_offset > 0 {
                    wide_character_offset -= 1;
                    continue;
                }

                let mut attributes = vec![
                    TermwizChange::AllAttributes(cell.attrs().clone()),
                    cell.str().into(),
                ];
                wide_character_offset = cell.width() - 1;

                changes.append(&mut attributes);
            }
        }

        self.cursor_state(&mut changes)?;

        Ok(changes)
    }

    /// Add the current cursor state.
    fn cursor_state(
        &self,
        changes: &mut Vec<TermwizChange>,
    ) -> Result<(), crate::errors::ShadowTerminalError> {
        let cursor = self.terminal.cursor_pos();

        let x = cursor.x;
        let y = cursor.y.try_into().with_whatever_context(|err| {
            format!("Couldn't convert cursor position to usize: {err:?}")
        })?;
        changes.push(TermwizChange::CursorPosition {
            x: TermwizPosition::Absolute(x),
            y: TermwizPosition::Absolute(y),
        });

        changes.push(TermwizChange::CursorShape(cursor.shape));
        changes.push(TermwizChange::CursorVisibility(cursor.visibility));

        Ok(())
    }

    /// Calculate the IDs of the lines that need to be output. Could just be the changed lines, or
    /// all the lines of the screen/scrollback.
    fn calculate_line_ids(
        &mut self,
        kind: &SurfaceKind,
        maybe_dirty_lines: Option<Vec<isize>>,
    ) -> Result<(Vec<usize>, usize), crate::errors::ShadowTerminalError> {
        let tty_size = self.terminal.get_size();
        let screen = self.terminal.screen_mut();
        let mut line_ids: Vec<usize> = Vec::new();
        let (output_start, output_end) = match kind {
            SurfaceKind::Scrollback => (0, screen.scrollback_rows()),
            SurfaceKind::Screen => {
                let end = screen.scrollback_rows() - self.scroll_position;
                let start = end - tty_size.rows;
                (start, end)
            }
        };

        match maybe_dirty_lines {
            Some(dirty_lines) => {
                for stable_dirty_line in dirty_lines {
                    let physical_line_id = screen
                        .stable_row_to_phys(stable_dirty_line)
                        .with_whatever_context(|| {
                            "Couldn't get physical row ID from stable row ID"
                        })?;
                    line_ids.push(physical_line_id);
                }
            }
            None => {
                for line_id in output_start..output_end {
                    line_ids.push(line_id);
                }
            }
        }

        Ok((line_ids, output_start))
    }
}

#[cfg(test)]
mod test {
    #[cfg(not(target_os = "windows"))]
    #[tokio::test(flavor = "multi_thread")]
    async fn wide_characters() {
        let mut stepper = Box::pin(crate::tests::helpers::run(Some(100), None)).await;
        let columns = stepper.shadow_terminal.terminal.get_size().cols;
        let full_row = "😀".repeat(columns.div_euclid(2));

        let command = format!("echo {full_row}");
        stepper.send_command(command.as_str()).unwrap();

        // The test code that dumps the terminal contents also includes the required blank cell(s)
        // that always follow wide characters. So we need to match them too.
        let raw_with_spaces = full_row
            .chars()
            .map(|character| character.to_string())
            .collect::<Vec<String>>()
            .join(" ");

        stepper
            .wait_for_string(&raw_with_spaces, None)
            .await
            .unwrap();
    }
}