zlayer_builder/tui/

app.rs

//! Main TUI application
//!
//! This module contains the main `BuildTui` application that manages
//! the terminal UI, event processing, and rendering loop.

use std::io::{self, Stdout};
use std::sync::mpsc::{Receiver, TryRecvError};

use crossterm::event::{self, Event, KeyCode, KeyEventKind, KeyModifiers};
use ratatui::prelude::*;
use ratatui::Terminal;
use zlayer_tui::terminal::{restore_terminal, setup_terminal, POLL_DURATION};
use zlayer_tui::widgets::scrollable_pane::OutputLine;

use super::build_view::BuildView;
use super::{BuildEvent, InstructionStatus, PlannedStage};

/// Main TUI application for build progress visualization
pub struct BuildTui {
    /// Channel to receive build events
    event_rx: Receiver<BuildEvent>,
    /// Current build state
    state: BuildState,
    /// Whether to keep running
    running: bool,
}

/// Internal build state tracking
#[derive(Debug, Default)]
pub struct BuildState {
    /// All stages in the build
    pub stages: Vec<StageState>,
    /// Current stage index
    pub current_stage: usize,
    /// Current instruction index within the current stage
    pub current_instruction: usize,
    /// Output lines from the build process
    pub output_lines: Vec<OutputLine>,
    /// Current scroll offset for output
    pub scroll_offset: usize,
    /// Whether the build has completed
    pub completed: bool,
    /// Error message if build failed
    pub error: Option<String>,
    /// Final image ID if build succeeded
    pub image_id: Option<String>,
    /// Total stages pre-declared by the backend via `BuildEvent::BuildStarted`.
    ///
    /// Populated once up-front. Used by `current_stage_display` so the
    /// denominator is stable rather than growing as stages arrive.
    /// `0` means no backend pre-declaration was received; fall back to
    /// `stages.len()` in that case.
    pub total_stages: usize,
    /// Total instructions pre-declared by the backend via
    /// `BuildEvent::BuildStarted`.
    ///
    /// Populated once up-front. Used by `total_instructions` so the
    /// progress-bar denominator is stable rather than growing as
    /// `InstructionStarted` events arrive. `0` means no backend
    /// pre-declaration was received; fall back to the event-sum.
    pub total_instructions: usize,
}

/// State of a single build stage
#[derive(Debug, Clone)]
pub struct StageState {
    /// Stage index
    pub index: usize,
    /// Optional stage name
    pub name: Option<String>,
    /// Base image for this stage
    pub base_image: String,
    /// Instructions in this stage
    pub instructions: Vec<InstructionState>,
    /// Whether the stage is complete
    pub complete: bool,
}

/// State of a single instruction
#[derive(Debug, Clone)]
pub struct InstructionState {
    /// Instruction text
    pub text: String,
    /// Current status
    pub status: InstructionStatus,
}

impl BuildTui {
    /// Create a new TUI with an event receiver
    #[must_use]
    pub fn new(event_rx: Receiver<BuildEvent>) -> Self {
        Self {
            event_rx,
            state: BuildState::default(),
            running: true,
        }
    }

    /// Run the TUI (blocking)
    ///
    /// This will take over the terminal, display the build progress,
    /// and return when the build completes or the user quits (q key).
    ///
    /// # Errors
    ///
    /// Returns an error if terminal setup, rendering, or restoration fails.
    pub fn run(&mut self) -> io::Result<()> {
        let mut terminal = setup_terminal()?;
        let result = self.run_loop(&mut terminal);
        restore_terminal(&mut terminal)?;
        result
    }

    /// Main event loop
    fn run_loop(&mut self, terminal: &mut Terminal<CrosstermBackend<Stdout>>) -> io::Result<()> {
        while self.running {
            // Process any pending build events
            self.process_events();

            // Render the UI
            terminal.draw(|frame| self.render(frame))?;

            // Auto-exit when build completes (success or failure).
            // The TUI disappears and the caller prints a summary to stdout.
            if self.state.completed {
                break;
            }

            // Handle keyboard input with a short timeout
            if event::poll(POLL_DURATION)? {
                if let Event::Key(key) = event::read()? {
                    if key.kind == KeyEventKind::Press {
                        self.handle_input(key.code, key.modifiers);
                    }
                }
            }
        }

        Ok(())
    }

    /// Process pending build events from the channel
    fn process_events(&mut self) {
        loop {
            match self.event_rx.try_recv() {
                Ok(event) => self.handle_build_event(event),
                Err(TryRecvError::Empty) => break,
                Err(TryRecvError::Disconnected) => {
                    // Build process has finished sending events
                    if !self.state.completed {
                        // If we didn't get a completion event, treat as unexpected end
                        self.state.completed = true;
                        if self.state.error.is_none() && self.state.image_id.is_none() {
                            self.state.error = Some("Build ended unexpectedly".to_string());
                        }
                    }
                    break;
                }
            }
        }
    }

    /// Handle a single build event
    fn handle_build_event(&mut self, event: BuildEvent) {
        match event {
            BuildEvent::BuildStarted {
                total_stages,
                total_instructions,
            } => {
                self.state.total_stages = total_stages;
                self.state.total_instructions = total_instructions;
            }

            BuildEvent::BuildPlan { stages } => {
                self.populate_from_plan(stages);
            }

            BuildEvent::StageStarted {
                index,
                name,
                base_image,
            } => {
                self.handle_stage_started(index, name, base_image);
            }

            BuildEvent::InstructionStarted {
                stage,
                index,
                instruction,
            } => {
                if let Some(stage_state) = self.state.stages.get_mut(stage) {
                    // Ensure we have enough instructions
                    while stage_state.instructions.len() <= index {
                        stage_state.instructions.push(InstructionState {
                            text: String::new(),
                            status: InstructionStatus::Pending,
                        });
                    }

                    // Update the instruction
                    stage_state.instructions[index] = InstructionState {
                        text: instruction,
                        status: InstructionStatus::Running,
                    };
                    self.state.current_instruction = index;
                }
            }

            BuildEvent::Output { line, is_stderr } => {
                self.state.output_lines.push(OutputLine {
                    text: line,
                    is_stderr,
                });

                // Auto-scroll to bottom if we were already at the bottom
                let visible_lines = 10; // Approximate
                let max_scroll = self.state.output_lines.len().saturating_sub(visible_lines);
                if self.state.scroll_offset >= max_scroll.saturating_sub(1) {
                    self.state.scroll_offset =
                        self.state.output_lines.len().saturating_sub(visible_lines);
                }
            }

            BuildEvent::InstructionComplete {
                stage,
                index,
                cached,
            } => {
                if let Some(stage_state) = self.state.stages.get_mut(stage) {
                    if let Some(inst) = stage_state.instructions.get_mut(index) {
                        inst.status = InstructionStatus::Complete { cached };
                    }
                }
            }

            BuildEvent::StageComplete { index } => {
                if let Some(stage_state) = self.state.stages.get_mut(index) {
                    stage_state.complete = true;
                }
            }

            BuildEvent::BuildComplete { image_id } => {
                self.state.completed = true;
                self.state.image_id = Some(image_id);
            }

            BuildEvent::BuildFailed { error } => {
                self.state.completed = true;
                self.state.error = Some(error);

                // Mark current instruction as failed
                if let Some(stage_state) = self.state.stages.get_mut(self.state.current_stage) {
                    if let Some(inst) = stage_state
                        .instructions
                        .get_mut(self.state.current_instruction)
                    {
                        if inst.status.is_running() {
                            inst.status = InstructionStatus::Failed;
                        }
                    }
                }
            }
        }
    }

    /// Handle a [`BuildEvent::StageStarted`]: (re)create the stage at `index`
    /// with an empty instruction list and make it current.
    ///
    /// Note: this resets the stage's instructions to empty, which would wipe a
    /// plan pre-filled by [`BuildEvent::BuildPlan`]. The buildah sidecar
    /// therefore does NOT emit `StageStarted` after a `BuildPlan`; only the
    /// native backend (which never sends `BuildPlan`) does.
    fn handle_stage_started(&mut self, index: usize, name: Option<String>, base_image: String) {
        // Ensure we have enough stages
        while self.state.stages.len() <= index {
            self.state.stages.push(StageState {
                index: self.state.stages.len(),
                name: None,
                base_image: String::new(),
                instructions: Vec::new(),
                complete: false,
            });
        }

        // Update the stage
        self.state.stages[index] = StageState {
            index,
            name,
            base_image,
            instructions: Vec::new(),
            complete: false,
        };
        self.state.current_stage = index;
        self.state.current_instruction = 0;
    }

    /// Pre-populate the state's stages/instructions from a [`BuildEvent::BuildPlan`].
    ///
    /// One [`StageState`] per planned stage (in order, named, with its base
    /// image and `complete: false`); each planned instruction becomes a
    /// `Pending` [`InstructionState`]. `total_stages`/`total_instructions`
    /// are seeded from the plan only if still `0`. After this, the sidecar
    /// will NOT send `StageStarted` (which would wipe a stage's instructions);
    /// it only sends `InstructionStarted`/`InstructionComplete`, whose
    /// existing handlers update `instructions[index]` in place.
    fn populate_from_plan(&mut self, stages: Vec<PlannedStage>) {
        let plan_stage_count = stages.len();
        let plan_instruction_count: usize = stages.iter().map(|s| s.instructions.len()).sum();

        self.state.stages = stages
            .into_iter()
            .enumerate()
            .map(|(index, planned)| StageState {
                index,
                name: planned.name,
                base_image: planned.base_image,
                instructions: planned
                    .instructions
                    .into_iter()
                    .map(|text| InstructionState {
                        text,
                        status: InstructionStatus::Pending,
                    })
                    .collect(),
                complete: false,
            })
            .collect();

        if self.state.total_stages == 0 {
            self.state.total_stages = plan_stage_count;
        }
        if self.state.total_instructions == 0 {
            self.state.total_instructions = plan_instruction_count;
        }

        self.state.current_stage = 0;
        self.state.current_instruction = 0;
    }

    /// Handle keyboard input
    fn handle_input(&mut self, key: KeyCode, modifiers: KeyModifiers) {
        match key {
            KeyCode::Char('q') | KeyCode::Esc => {
                self.running = false;
            }
            KeyCode::Char('c') if modifiers.contains(KeyModifiers::CONTROL) => {
                self.running = false;
            }
            KeyCode::Up | KeyCode::Char('k') => {
                self.state.scroll_offset = self.state.scroll_offset.saturating_sub(1);
            }
            KeyCode::Down | KeyCode::Char('j') => {
                let max_scroll = self.state.output_lines.len().saturating_sub(10);
                if self.state.scroll_offset < max_scroll {
                    self.state.scroll_offset += 1;
                }
            }
            KeyCode::PageUp => {
                self.state.scroll_offset = self.state.scroll_offset.saturating_sub(10);
            }
            KeyCode::PageDown => {
                let max_scroll = self.state.output_lines.len().saturating_sub(10);
                self.state.scroll_offset = (self.state.scroll_offset + 10).min(max_scroll);
            }
            KeyCode::Home => {
                self.state.scroll_offset = 0;
            }
            KeyCode::End => {
                let max_scroll = self.state.output_lines.len().saturating_sub(10);
                self.state.scroll_offset = max_scroll;
            }
            _ => {}
        }
    }

    /// Render the UI
    fn render(&self, frame: &mut Frame) {
        let view = BuildView::new(&self.state);
        frame.render_widget(view, frame.area());
    }
}

impl BuildState {
    /// Get the total number of instructions across all stages.
    ///
    /// Prefers the pre-declared total from `BuildEvent::BuildStarted`
    /// (so the progress-bar denominator is stable); falls back to
    /// summing the stages we've learned about so far for older tests
    /// or backends that don't emit `BuildStarted`.
    #[must_use]
    pub fn total_instructions(&self) -> usize {
        let event_sum: usize = self.stages.iter().map(|s| s.instructions.len()).sum();
        self.total_instructions.max(event_sum)
    }

    /// Get the number of completed instructions across all stages
    #[must_use]
    pub fn completed_instructions(&self) -> usize {
        self.stages
            .iter()
            .flat_map(|s| s.instructions.iter())
            .filter(|i| i.status.is_complete())
            .count()
    }

    /// Get a display string for the current stage
    #[must_use]
    pub fn current_stage_display(&self) -> String {
        if let Some(stage) = self.stages.get(self.current_stage) {
            let name_part = stage
                .name
                .as_ref()
                .map(|n| format!("{n} "))
                .unwrap_or_default();
            format!(
                "Stage {}/{}: {}({})",
                self.current_stage + 1,
                self.total_stages.max(self.stages.len()).max(1),
                name_part,
                stage.base_image
            )
        } else {
            "Initializing...".to_string()
        }
    }
}

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

    #[test]
    fn test_build_state_default() {
        let state = BuildState::default();
        assert!(state.stages.is_empty());
        assert!(!state.completed);
        assert!(state.error.is_none());
        assert!(state.image_id.is_none());
    }

    #[test]
    fn test_build_state_instruction_counts() {
        let mut state = BuildState::default();
        state.stages.push(StageState {
            index: 0,
            name: None,
            base_image: "alpine".to_string(),
            instructions: vec![
                InstructionState {
                    text: "RUN echo 1".to_string(),
                    status: InstructionStatus::Complete { cached: false },
                },
                InstructionState {
                    text: "RUN echo 2".to_string(),
                    status: InstructionStatus::Running,
                },
            ],
            complete: false,
        });

        assert_eq!(state.total_instructions(), 2);
        assert_eq!(state.completed_instructions(), 1);
    }

    #[test]
    fn total_instructions_uses_predeclared_total() {
        // Reproduces the old bug: before this fix, `total_instructions()`
        // summed stages as `InstructionStarted` events arrived, so the
        // TUI denominator grew in lockstep with the numerator. With the
        // up-front `BuildStarted` event, the denominator is stable from
        // the first instruction onward.
        let (tx, rx) = mpsc::channel();
        let mut tui = BuildTui::new(rx);

        tx.send(BuildEvent::BuildStarted {
            total_stages: 2,
            total_instructions: 7,
        })
        .unwrap();

        tx.send(BuildEvent::StageStarted {
            index: 0,
            name: None,
            base_image: "alpine".to_string(),
        })
        .unwrap();

        tx.send(BuildEvent::InstructionStarted {
            stage: 0,
            index: 0,
            instruction: "RUN foo".to_string(),
        })
        .unwrap();

        drop(tx);
        tui.process_events();

        // Pre-declared total wins over the in-progress event-sum of 1.
        assert_eq!(tui.state.total_instructions(), 7);
        // Current-stage denominator is also populated up-front.
        assert!(tui.state.current_stage_display().contains("Stage 1/2"));
    }

    #[test]
    fn build_plan_prefills_pending_instructions() {
        let (tx, rx) = mpsc::channel();
        let mut tui = BuildTui::new(rx);

        tx.send(BuildEvent::BuildPlan {
            stages: vec![
                PlannedStage {
                    name: Some("builder".to_string()),
                    base_image: "node:20-alpine".to_string(),
                    instructions: vec!["WORKDIR /app".to_string(), "RUN npm ci".to_string()],
                },
                PlannedStage {
                    name: None,
                    base_image: "alpine".to_string(),
                    instructions: vec!["COPY --from=builder /app /app".to_string()],
                },
            ],
        })
        .unwrap();

        drop(tx);
        tui.process_events();

        // Full plan is present immediately — two stages, three instructions,
        // all Pending. (This is what unblocks "Waiting for build to start...".)
        assert_eq!(tui.state.stages.len(), 2);
        assert_eq!(tui.state.total_stages, 2);
        assert_eq!(tui.state.total_instructions(), 3);
        assert_eq!(tui.state.completed_instructions(), 0);
        assert!(tui
            .state
            .stages
            .iter()
            .flat_map(|s| &s.instructions)
            .all(|i| matches!(i.status, InstructionStatus::Pending)));
        assert_eq!(tui.state.stages[0].instructions[0].text, "WORKDIR /app");
    }

    #[test]
    fn build_plan_then_progress_advances_statuses() {
        // Mirrors what the sidecar backend emits: a BuildPlan, then a
        // sequence of InstructionStarted/InstructionComplete derived from
        // buildah's `--> ` commit markers. Statuses must advance
        // Pending -> Running -> Complete and the counter must track.
        let (tx, rx) = mpsc::channel();
        let mut tui = BuildTui::new(rx);

        tx.send(BuildEvent::BuildPlan {
            stages: vec![PlannedStage {
                name: None,
                base_image: "alpine".to_string(),
                instructions: vec!["RUN echo a".to_string(), "RUN echo b".to_string()],
            }],
        })
        .unwrap();

        // Sidecar marks the first instruction Running up-front.
        tx.send(BuildEvent::InstructionStarted {
            stage: 0,
            index: 0,
            instruction: "RUN echo a".to_string(),
        })
        .unwrap();

        drop(tx);
        tui.process_events();

        assert!(tui.state.stages[0].instructions[0].status.is_running());
        assert!(matches!(
            tui.state.stages[0].instructions[1].status,
            InstructionStatus::Pending
        ));
        assert_eq!(tui.state.completed_instructions(), 0);

        // First `--> ` marker: complete inst 0, start inst 1.
        let (tx, rx) = mpsc::channel();
        let mut tui2 = BuildTui::new(rx);
        tx.send(BuildEvent::BuildPlan {
            stages: vec![PlannedStage {
                name: None,
                base_image: "alpine".to_string(),
                instructions: vec!["RUN echo a".to_string(), "RUN echo b".to_string()],
            }],
        })
        .unwrap();
        tx.send(BuildEvent::InstructionStarted {
            stage: 0,
            index: 0,
            instruction: "RUN echo a".to_string(),
        })
        .unwrap();
        tx.send(BuildEvent::InstructionComplete {
            stage: 0,
            index: 0,
            cached: false,
        })
        .unwrap();
        tx.send(BuildEvent::InstructionStarted {
            stage: 0,
            index: 1,
            instruction: "RUN echo b".to_string(),
        })
        .unwrap();
        tx.send(BuildEvent::InstructionComplete {
            stage: 0,
            index: 1,
            cached: true,
        })
        .unwrap();
        drop(tx);
        tui2.process_events();

        assert!(tui2.state.stages[0].instructions[0].status.is_complete());
        assert!(tui2.state.stages[0].instructions[1].status.is_complete());
        assert_eq!(tui2.state.completed_instructions(), 2);
        assert_eq!(tui2.state.total_instructions(), 2);
    }

    #[test]
    fn test_handle_stage_started() {
        let (tx, rx) = mpsc::channel();
        let mut tui = BuildTui::new(rx);

        tx.send(BuildEvent::StageStarted {
            index: 0,
            name: Some("builder".to_string()),
            base_image: "node:20".to_string(),
        })
        .unwrap();

        drop(tx);
        tui.process_events();

        assert_eq!(tui.state.stages.len(), 1);
        assert_eq!(tui.state.stages[0].name, Some("builder".to_string()));
        assert_eq!(tui.state.stages[0].base_image, "node:20");
    }

    #[test]
    fn test_handle_instruction_lifecycle() {
        let (tx, rx) = mpsc::channel();
        let mut tui = BuildTui::new(rx);

        // Start stage
        tx.send(BuildEvent::StageStarted {
            index: 0,
            name: None,
            base_image: "alpine".to_string(),
        })
        .unwrap();

        // Start instruction
        tx.send(BuildEvent::InstructionStarted {
            stage: 0,
            index: 0,
            instruction: "RUN echo hello".to_string(),
        })
        .unwrap();

        // Complete instruction
        tx.send(BuildEvent::InstructionComplete {
            stage: 0,
            index: 0,
            cached: true,
        })
        .unwrap();

        drop(tx);
        tui.process_events();

        let inst = &tui.state.stages[0].instructions[0];
        assert_eq!(inst.text, "RUN echo hello");
        assert!(matches!(
            inst.status,
            InstructionStatus::Complete { cached: true }
        ));
    }

    #[test]
    fn test_handle_build_complete() {
        let (tx, rx) = mpsc::channel();
        let mut tui = BuildTui::new(rx);

        tx.send(BuildEvent::BuildComplete {
            image_id: "sha256:abc123".to_string(),
        })
        .unwrap();

        drop(tx);
        tui.process_events();

        assert!(tui.state.completed);
        assert_eq!(tui.state.image_id, Some("sha256:abc123".to_string()));
        assert!(tui.state.error.is_none());
    }

    #[test]
    fn test_handle_build_failed() {
        let (tx, rx) = mpsc::channel();
        let mut tui = BuildTui::new(rx);

        // Set up a running instruction first
        tx.send(BuildEvent::StageStarted {
            index: 0,
            name: None,
            base_image: "alpine".to_string(),
        })
        .unwrap();

        tx.send(BuildEvent::InstructionStarted {
            stage: 0,
            index: 0,
            instruction: "RUN exit 1".to_string(),
        })
        .unwrap();

        tx.send(BuildEvent::BuildFailed {
            error: "Command failed with exit code 1".to_string(),
        })
        .unwrap();

        drop(tx);
        tui.process_events();

        assert!(tui.state.completed);
        assert!(tui.state.error.is_some());
        assert!(tui.state.stages[0].instructions[0].status.is_failed());
    }
}