Skip to main content

intelli_shell/
tui.rs

1use std::{
2    cmp,
3    io::{self, Stdout, stdout},
4    ops::{Deref, DerefMut},
5    thread,
6    time::Duration,
7};
8
9use color_eyre::Result;
10use crossterm::{
11    cursor,
12    event::{
13        self, Event as CrosstermEvent, EventStream, KeyCode, KeyEvent, KeyEventKind, KeyModifiers,
14        KeyboardEnhancementFlags, MouseEvent,
15    },
16    style,
17    terminal::{self, ClearType, supports_keyboard_enhancement},
18};
19use futures_util::{FutureExt, StreamExt};
20use ratatui::{CompletedFrame, Frame, Terminal, backend::CrosstermBackend as Backend, layout::Rect};
21use serde::{Deserialize, Serialize};
22use tokio::{
23    sync::mpsc::{self, UnboundedReceiver, UnboundedSender},
24    task::JoinHandle,
25    time::interval,
26};
27use tokio_util::sync::CancellationToken;
28use tracing::instrument;
29
30/// Events that can occur within the TUI application
31#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
32pub enum Event {
33    /// A periodic tick event, useful for time-based updates or animations
34    Tick,
35    /// A periodic render event, suggesting the UI should be redrawn
36    Render,
37    /// The terminal window gained focus
38    FocusGained,
39    /// The terminal window lost focus
40    FocusLost,
41    /// Text was pasted into the terminal (requires paste mode)
42    Paste(String),
43    /// A key was pressed
44    Key(KeyEvent),
45    /// A mouse event occurred (requires mouse capture)
46    Mouse(MouseEvent),
47    /// The terminal window was resized (columns and rows)
48    Resize(u16, u16),
49}
50
51/// Manages the terminal User Interface (TUI) lifecycle, event handling, and rendering
52pub struct Tui {
53    stdout: Stdout,
54    terminal: Terminal<Backend<Stdout>>,
55    task: JoinHandle<()>,
56    loop_cancellation_token: CancellationToken,
57    global_cancellation_token: CancellationToken,
58    event_rx: UnboundedReceiver<Event>,
59    event_tx: UnboundedSender<Event>,
60    frame_rate: f64,
61    tick_rate: f64,
62    mouse: bool,
63    paste: bool,
64    keyboard_enhancement: bool,
65    state: Option<State>,
66}
67
68#[derive(Clone, Copy)]
69enum State {
70    FullScreen(bool),
71    Inline(bool, InlineTuiContext),
72}
73
74#[derive(Clone, Copy)]
75struct InlineTuiContext {
76    min_height: u16,
77    x: u16,
78    y: u16,
79    restore_cursor_x: u16,
80    restore_cursor_y: u16,
81}
82
83#[allow(dead_code, reason = "provide a useful interface, even if not required yet")]
84impl Tui {
85    /// Constructs a new terminal ui with default settings
86    pub fn new(cancellation_token: CancellationToken) -> Result<Self> {
87        let (event_tx, event_rx) = mpsc::unbounded_channel();
88        Ok(Self {
89            stdout: stdout(),
90            terminal: Terminal::new(Backend::new(stdout()))?,
91            task: tokio::spawn(async {}),
92            loop_cancellation_token: CancellationToken::new(),
93            global_cancellation_token: cancellation_token,
94            event_rx,
95            event_tx,
96            frame_rate: 60.0,
97            tick_rate: 10.0,
98            mouse: false,
99            paste: false,
100            keyboard_enhancement: false,
101            state: None,
102        })
103    }
104
105    /// Sets the tick rate for the TUI.
106    ///
107    /// The tick rate determines how frequently `Event::Tick` is generated.
108    pub fn tick_rate(mut self, tick_rate: f64) -> Self {
109        self.state.is_some().then(|| panic!("Can't updated an entered TUI"));
110        self.tick_rate = tick_rate;
111        self
112    }
113
114    /// Sets the frame rate for the TUI.
115    ///
116    /// The frame rate determines how often `Event::Render` is emitted.
117    pub fn frame_rate(mut self, frame_rate: f64) -> Self {
118        self.state.is_some().then(|| panic!("Can't updated an entered TUI"));
119        self.frame_rate = frame_rate;
120        self
121    }
122
123    /// Enables or disables mouse event capture.
124    ///
125    /// If true, `Event::Mouse` events will be emitted.
126    pub fn mouse(mut self, mouse: bool) -> Self {
127        self.state.is_some().then(|| panic!("Can't updated an entered TUI"));
128        self.mouse = mouse;
129        self
130    }
131
132    /// Enables or disables bracketed paste mode.
133    ///
134    /// If true, `Event::Paste` events will be emitted.
135    pub fn paste(mut self, paste: bool) -> Self {
136        self.state.is_some().then(|| panic!("Can't updated an entered TUI"));
137        self.paste = paste;
138        self
139    }
140
141    /// Enables or disables keyboard enhancement flags.
142    pub fn keyboard_enhancement(mut self, keyboard_enhancement: bool) -> Self {
143        self.state.is_some().then(|| panic!("Can't updated an entered TUI"));
144        self.keyboard_enhancement = keyboard_enhancement;
145        self
146    }
147
148    /// Asynchronously retrieves the next event from the event queue.
149    ///
150    /// Returns `None` if the event channel has been closed (e.g., the event loop has stopped).
151    pub async fn next_event(&mut self) -> Option<Event> {
152        self.event_rx.recv().await
153    }
154
155    /// Prepares the terminal for full-screen TUI interaction and starts the event loop
156    pub fn enter(&mut self) -> Result<()> {
157        self.state.is_some().then(|| panic!("Can't re-enter on a TUI"));
158
159        tracing::trace!(mouse = self.mouse, paste = self.paste, "Entering a full-screen TUI");
160
161        // Enter raw mode and set up the terminal
162        let keyboard_enhancement_enabled = self.enter_raw_mode(true)?;
163
164        // Store the state and start the event loop
165        self.state = Some(State::FullScreen(keyboard_enhancement_enabled));
166        self.start();
167
168        Ok(())
169    }
170
171    /// Prepares the terminal for inline TUI interaction and starts the event loop
172    pub fn enter_inline(&mut self, extra_line: bool, min_height: u16) -> Result<()> {
173        self.state.is_some().then(|| panic!("Can't re-enter on a TUI"));
174        let extra_line = extra_line as u16;
175
176        tracing::trace!(
177            mouse = self.mouse,
178            paste = self.paste,
179            extra_line,
180            min_height,
181            "Entering an inline TUI"
182        );
183
184        // Save the original cursor position
185        let (orig_cursor_x, orig_cursor_y) = cursor::position()?;
186        tracing::trace!("Initial cursor position: ({orig_cursor_x},{orig_cursor_y})");
187        // Prepare the area for the inline content
188        crossterm::execute!(
189            self.stdout,
190            // Fill in the minimum height (plus the extra line), the cursor will end up at the end
191            style::Print("\n".repeat((min_height + extra_line) as usize)),
192            // Move the cursor back the min height (without the extra lines)
193            cursor::MoveToPreviousLine(min_height),
194            // And clear the lines below
195            terminal::Clear(ClearType::FromCursorDown)
196        )?;
197        // Retrieve the new cursor position, which defines the starting coords for the area to render in
198        let (cursor_x, cursor_y) = cursor::position()?;
199        // Calculate where the cursor should be restored to
200        let restore_cursor_x = orig_cursor_x;
201        let restore_cursor_y = cmp::min(orig_cursor_y, cmp::max(cursor_y, extra_line) - extra_line);
202        tracing::trace!("Cursor shall be restored at: ({restore_cursor_x},{restore_cursor_y})");
203
204        // Enter raw mode and set up the terminal
205        let keyboard_enhancement_enabled = self.enter_raw_mode(false)?;
206
207        // Store the state and start the event loop
208        self.state = Some(State::Inline(
209            keyboard_enhancement_enabled,
210            InlineTuiContext {
211                min_height,
212                x: cursor_x,
213                y: cursor_y,
214                restore_cursor_x,
215                restore_cursor_y,
216            },
217        ));
218        self.start();
219
220        Ok(())
221    }
222
223    /// Renders the TUI using the provided callback function.
224    ///
225    /// The callback receives a mutable reference to the `Frame` and the area to render in, which might not be the same
226    /// as the frame area for inline TUIs.
227    pub fn render<F>(&mut self, render_callback: F) -> io::Result<CompletedFrame<'_>>
228    where
229        F: FnOnce(&mut Frame, Rect),
230    {
231        let Some(state) = self.state else {
232            return Err(io::Error::other("Cannot render on a non-entered TUI"));
233        };
234
235        self.terminal.draw(|frame| {
236            let area = match state {
237                State::FullScreen(_) => frame.area(),
238                State::Inline(_, inline) => inline_render_area(frame.area(), inline),
239            };
240
241            render_callback(frame, area);
242        })
243    }
244
245    /// Restores the terminal to its original state and stops the event loop
246    pub fn exit(mut self) -> Result<()> {
247        self.state.is_none().then(|| panic!("Cannot exit a non-entered TUI"));
248        self.stop();
249        self.restore_terminal()
250    }
251
252    fn restore_terminal(&mut self) -> Result<()> {
253        match self.state.take() {
254            None => (),
255            Some(State::FullScreen(keyboard_enhancement_enabled)) => {
256                tracing::trace!("Leaving the full-screen TUI");
257                self.flush()?;
258                self.exit_raw_mode(true, keyboard_enhancement_enabled)?;
259            }
260            Some(State::Inline(keyboard_enhancement_enabled, ctx)) => {
261                tracing::trace!("Leaving the inline TUI");
262                self.flush()?;
263                self.exit_raw_mode(false, keyboard_enhancement_enabled)?;
264                crossterm::execute!(
265                    self.stdout,
266                    cursor::MoveTo(ctx.restore_cursor_x, ctx.restore_cursor_y),
267                    terminal::Clear(ClearType::FromCursorDown)
268                )?;
269            }
270        }
271
272        Ok(())
273    }
274
275    fn enter_raw_mode(&mut self, alt_screen: bool) -> Result<bool> {
276        terminal::enable_raw_mode()?;
277        crossterm::execute!(self.stdout, cursor::Hide)?;
278        if alt_screen {
279            crossterm::execute!(self.stdout, terminal::EnterAlternateScreen)?;
280        }
281        if self.mouse {
282            crossterm::execute!(self.stdout, event::EnableMouseCapture)?;
283        }
284        if self.paste {
285            crossterm::execute!(self.stdout, event::EnableBracketedPaste)?;
286        }
287
288        let mut ke_enabled = false;
289        if self.keyboard_enhancement {
290            tracing::trace!("Checking keyboard enhancement support");
291            let keyboard_enhancement_supported = supports_keyboard_enhancement()
292                .inspect_err(|err| tracing::error!("{err}"))
293                .unwrap_or(false);
294
295            if keyboard_enhancement_supported {
296                tracing::trace!("Keyboard enhancement flags enabled");
297                crossterm::execute!(
298                    self.stdout,
299                    event::PushKeyboardEnhancementFlags(
300                        KeyboardEnhancementFlags::DISAMBIGUATE_ESCAPE_CODES
301                            | KeyboardEnhancementFlags::REPORT_ALL_KEYS_AS_ESCAPE_CODES
302                            | KeyboardEnhancementFlags::REPORT_ALTERNATE_KEYS
303                    ),
304                )?;
305                ke_enabled = true;
306            } else {
307                tracing::debug!("Keyboard enhancement not supported");
308            }
309        } else {
310            tracing::trace!("Keyboard enhancement flags not enabled");
311        }
312
313        Ok(ke_enabled)
314    }
315
316    fn exit_raw_mode(&mut self, alt_screen: bool, keyboard_enhancement_enabled: bool) -> Result<()> {
317        if keyboard_enhancement_enabled {
318            crossterm::execute!(self.stdout, event::PopKeyboardEnhancementFlags)?;
319        }
320
321        if self.paste {
322            crossterm::execute!(self.stdout, event::DisableBracketedPaste)?;
323        }
324        if self.mouse {
325            crossterm::execute!(self.stdout, event::DisableMouseCapture)?;
326        }
327        if alt_screen {
328            crossterm::execute!(self.stdout, terminal::LeaveAlternateScreen)?;
329        }
330        crossterm::execute!(self.stdout, cursor::Show)?;
331        terminal::disable_raw_mode()?;
332
333        Ok(())
334    }
335
336    fn start(&mut self) {
337        self.cancel();
338        self.loop_cancellation_token = CancellationToken::new();
339
340        tracing::trace!(
341            tick_rate = self.tick_rate,
342            frame_rate = self.frame_rate,
343            "Starting the event loop"
344        );
345
346        self.task = tokio::spawn(Self::event_loop(
347            self.event_tx.clone(),
348            self.loop_cancellation_token.clone(),
349            self.global_cancellation_token.clone(),
350            self.tick_rate,
351            self.frame_rate,
352        ));
353    }
354
355    #[instrument(skip_all)]
356    async fn event_loop(
357        event_tx: UnboundedSender<Event>,
358        loop_cancellation_token: CancellationToken,
359        global_cancellation_token: CancellationToken,
360        tick_rate: f64,
361        frame_rate: f64,
362    ) {
363        let mut event_stream = EventStream::new();
364        let mut tick_interval = interval(Duration::from_secs_f64(1.0 / tick_rate));
365        let mut render_interval = interval(Duration::from_secs_f64(1.0 / frame_rate));
366
367        loop {
368            let event = tokio::select! {
369                // Ensure signals are checked in order (cancellation first)
370                biased;
371
372                // Exit the loop if any cancellation is requested
373                _ = loop_cancellation_token.cancelled() => {
374                    break;
375                }
376                _ = global_cancellation_token.cancelled() => {
377                    break;
378                }
379
380                // Crossterm events
381                crossterm_event = event_stream.next().fuse() => match crossterm_event {
382                    Some(Ok(event)) => match event {
383                        // On raw mode, SIGINT is no longer received and we should handle it manually
384                        CrosstermEvent::Key(KeyEvent {
385                            code: KeyCode::Char('c'),
386                            modifiers: KeyModifiers::CONTROL,
387                            ..
388                        }) => {
389                            tracing::debug!("Ctrl+C key event received in TUI, cancelling token");
390                            global_cancellation_token.cancel();
391                            continue;
392                        }
393                        // Process only key press events to avoid duplicate events for release/repeat
394                        CrosstermEvent::Key(key) if key.kind == KeyEventKind::Press => Event::Key(key),
395                        CrosstermEvent::Mouse(mouse) => Event::Mouse(mouse),
396                        CrosstermEvent::Resize(cols, rows) => Event::Resize(cols, rows),
397                        CrosstermEvent::FocusLost => Event::FocusLost,
398                        CrosstermEvent::FocusGained => Event::FocusGained,
399                        CrosstermEvent::Paste(s) => Event::Paste(s),
400                        _ => continue, // Ignore other crossterm event types
401                    }
402                    Some(Err(err)) =>  {
403                        tracing::error!("Error retrieving next crossterm event: {err}");
404                        break;
405                    },
406                    None => break, // Event stream ended, exit the loop
407                },
408
409                // Intervals
410                _ = tick_interval.tick() => Event::Tick,
411                _ = render_interval.tick() => Event::Render,
412            };
413
414            // Try to send the processed event
415            if event_tx.send(event).is_err() {
416                // If sending fails, the receiver is likely dropped. Exit the loop
417                break;
418            }
419        }
420
421        // Ensure the token is cancelled if the loop exits for reasons other than direct cancellation
422        // (e.g. event_stream ending or send error).
423        loop_cancellation_token.cancel();
424    }
425
426    fn stop(&self) {
427        if !self.task.is_finished() {
428            tracing::trace!("Stopping the event loop");
429            self.cancel();
430            let mut counter = 0;
431            while !self.task.is_finished() {
432                thread::sleep(Duration::from_millis(1));
433                counter += 1;
434                // Attempt to abort the task if it hasn't finished in a short period
435                if counter > 50 {
436                    tracing::debug!("Task hasn't finished in 50 milliseconds, attempting to abort");
437                    self.task.abort();
438                }
439                // Log an error and give up waiting if the task hasn't finished after the abort
440                if counter > 100 {
441                    tracing::error!("Failed to abort task in 100 milliseconds for unknown reason");
442                    break;
443                }
444            }
445        }
446    }
447
448    fn cancel(&self) {
449        self.loop_cancellation_token.cancel();
450    }
451}
452
453fn inline_render_area(frame: Rect, inline: InlineTuiContext) -> Rect {
454    if frame.width == 0 || frame.height == 0 {
455        return frame;
456    }
457    let min_height = inline.min_height.max(1).min(frame.height);
458    let x = inline.x.clamp(frame.x, frame.right() - 1);
459    let y = inline.y.clamp(frame.y, frame.bottom() - min_height);
460    Rect::new(x, y, frame.right() - x, frame.bottom() - y)
461}
462
463impl Deref for Tui {
464    type Target = Terminal<Backend<Stdout>>;
465
466    fn deref(&self) -> &Self::Target {
467        &self.terminal
468    }
469}
470
471impl DerefMut for Tui {
472    fn deref_mut(&mut self) -> &mut Self::Target {
473        &mut self.terminal
474    }
475}
476
477impl Drop for Tui {
478    fn drop(&mut self) {
479        self.stop();
480        if let Err(err) = self.restore_terminal() {
481            tracing::error!("Failed to restore terminal state: {err:?}");
482        }
483    }
484}
485
486#[cfg(test)]
487mod tests {
488    use ratatui::layout::Rect;
489
490    use super::{InlineTuiContext, inline_render_area};
491
492    #[test]
493    fn inline_area_stays_inside_a_shrunken_terminal() {
494        let frame = Rect::new(0, 0, 8, 4);
495        let inline = InlineTuiContext {
496            min_height: 6,
497            x: 12,
498            y: 6,
499            restore_cursor_x: 0,
500            restore_cursor_y: 0,
501        };
502
503        assert_eq!(inline_render_area(frame, inline), Rect::new(7, 0, 1, 4));
504        assert_eq!(inline_render_area(Rect::new(0, 0, 0, 0), inline), Rect::new(0, 0, 0, 0));
505    }
506
507    #[test]
508    fn inline_area_fits_minimum_height_and_respects_frame_origin() {
509        let inline = InlineTuiContext {
510            min_height: 3,
511            x: 12,
512            y: 10,
513            restore_cursor_x: 0,
514            restore_cursor_y: 0,
515        };
516        assert_eq!(inline_render_area(Rect::new(2, 3, 8, 4), inline), Rect::new(9, 4, 1, 3));
517    }
518}