galaxy/app/src/terminal/wormhole/trigger_state.rs

use std::collections::HashMap;
use std::sync::Arc;

use galaxy_core::ui::appearance::Appearance;
use galaxyui::r#async::SpawnedFutureHandle;
use galaxyui::{EntityId, SingletonEntity as _, ViewContext, ViewHandle};
use parking_lot::FairMutex;

use super::success_block::WormholeSuccessBlock;
use crate::terminal::model::block::BlockId;
use crate::terminal::model::session::SessionId;
use crate::terminal::model::terminal_model::SubshellInitializationInfo;
use crate::terminal::settings::TerminalSettings;
use crate::terminal::shell::ShellType;
use crate::terminal::{TerminalModel, TerminalView};

/// A unique identifier for a subshell separator.
pub type SeparatorId = usize;

/// These are elements in the BlockList which are similar to inline banners but are smaller, and
/// only meant to render in compact mode when their in-padding flag counterparts don't have enough
/// space in the padding to render.
#[derive(Default)]
struct SubshellSeparatorState {
    /// The ID for the next separator to be created.
    next_separator_id: SeparatorId,

    /// These are for rendering above the first block of a subshell session.
    separators: HashMap<SeparatorId, String>,
}

impl SubshellSeparatorState {
    /// Returns the ID to assign to the next separator
    fn next_separator_id(&mut self) -> SeparatorId {
        let next_id = self.next_separator_id;
        self.next_separator_id += 1;
        next_id
    }
}

#[derive(Debug)]
pub enum SshBlockState {
    WormholeSuccess {
        handle: ViewHandle<WormholeSuccessBlock>,
    },
}

impl SshBlockState {
    pub fn should_prevent_input(&self) -> bool {
        true
    }

    pub fn get_block_view_id(&self) -> EntityId {
        match self {
            SshBlockState::WormholeSuccess { handle, .. } => handle.id(),
        }
    }

    pub fn on_wormholed_session_complete(
        &self,
        ctx: &mut ViewContext<TerminalView>,
    ) -> Option<EntityId> {
        match self {
            SshBlockState::WormholeSuccess { handle } => {
                handle.update(ctx, |block, ctx| {
                    block.on_wormholed_session_complete(ctx);
                });
            }
        }
        None
    }
}

/// Temporary state used to trigger Wormholing.
#[derive(Default)]
struct WormholeTriggerState {
    block_id: Option<BlockId>,

    /// Lets us abort an attempt to auto wormhole if the subshell command
    /// hasn't completed.
    auto_wormhole_abort_handle: Option<SpawnedFutureHandle>,

    /// The subshell banner waits 1s before showing. This is to see that the command stays running
    /// for a while without exiting. We store the abort handle here so that the
    /// TerminalEvent::BlockCompleted event can abort the banner.
    subshell_banner_abort_handle: Option<SpawnedFutureHandle>,

    /// The command which may trigger ssh Wormholing
    pending_command: Option<String>,
    /// The Host which may trigger ssh Wormholing
    pending_wormhole_ssh_host: Option<String>,

    /// Which, if any, SSH block is currently added to the blocklist.
    ssh_block_state: Option<SshBlockState>,

    ssh_wormhole_timeout_handle: Option<SpawnedFutureHandle>,

    shell_type: Option<ShellType>,

    is_shell_detection_in_progress: bool,
}

#[derive(Default)]
pub struct WormholeState {
    session_id: Option<SessionId>,

    pending_state: Option<WormholeTriggerState>,
    /// Stores the metadata needed to render any separators above the first block of a subshell.
    subshell_separator_state: SubshellSeparatorState,
    /// A unique-enough ID that is used to validate that a timeout is still valid.
    timeout_id: u8,
}

impl WormholeState {
    pub fn delete_state(&mut self) {
        self.pending_state.take();
    }

    pub fn is_shell_detection_in_progress(&self) -> bool {
        self.pending_state
            .as_ref()
            .map(|state| state.is_shell_detection_in_progress)
            .unwrap_or_default()
    }

    pub fn set_shell_detection_in_progress(&mut self) {
        if let Some(ref mut pending_state) = self.pending_state.as_mut() {
            pending_state.is_shell_detection_in_progress = true;
        }
    }

    pub fn set_shell_type(&mut self, shell_type: &ShellType) {
        let pending_state = self.pending_state.get_or_insert_with(Default::default);
        pending_state.shell_type = Some(shell_type.to_owned());
        pending_state.is_shell_detection_in_progress = false;
    }

    pub fn get_shell_type(&self) -> Option<ShellType> {
        self.pending_state
            .as_ref()
            .and_then(|state| state.shell_type)
    }

    pub fn add_subshell_separator(
        &mut self,
        subshell_info: &SubshellInitializationInfo,
        terminal_model: Arc<FairMutex<TerminalModel>>,
        ctx: &mut ViewContext<TerminalView>,
    ) {
        let Some(command) = subshell_info.spawning_command.split_whitespace().next() else {
            return;
        };
        let separator_id = self.subshell_separator_state.next_separator_id();
        let appearance = Appearance::as_ref(ctx);
        let terminal_spacing =
            TerminalSettings::as_ref(ctx).terminal_spacing(appearance.line_height_ratio(), ctx);
        let height = terminal_spacing.subshell_separator_height;
        self.subshell_separator_state
            .separators
            .insert(separator_id, command.to_owned());
        terminal_model
            .lock()
            .block_list_mut()
            .append_subshell_separator(separator_id, height);
        ctx.notify();
    }

    pub fn get_subshell_separators(&self) -> &HashMap<SeparatorId, String> {
        &self.subshell_separator_state.separators
    }

    pub fn add_subshell_banner_abort_handle(&mut self, spawned_future_handle: SpawnedFutureHandle) {
        let pending_state = self.pending_state.get_or_insert_with(Default::default);
        pending_state.subshell_banner_abort_handle = Some(spawned_future_handle);
    }

    pub fn take_subshell_banner_abort_handle(&mut self) -> Option<SpawnedFutureHandle> {
        self.pending_state
            .as_mut()
            .and_then(|state| state.subshell_banner_abort_handle.take())
    }

    pub fn add_auto_wormhole_abort_handle(&mut self, spawned_future_handle: SpawnedFutureHandle) {
        let pending_state = self.pending_state.get_or_insert_with(Default::default);
        pending_state.auto_wormhole_abort_handle = Some(spawned_future_handle);
    }

    pub fn abort_auto_wormhole(&mut self) {
        if let Some(abort_handle) = self
            .pending_state
            .as_mut()
            .and_then(|state| state.auto_wormhole_abort_handle.take())
        {
            abort_handle.abort();
        };
    }

    pub fn add_ssh_wormhole_timeout_handle(&mut self, spawned_future_handle: SpawnedFutureHandle) {
        let pending_state = self.pending_state.get_or_insert_with(Default::default);
        pending_state.ssh_wormhole_timeout_handle = Some(spawned_future_handle);
    }

    pub fn abort_ssh_wormhole_timeout(&mut self) {
        self.replace_timeout_id();
        if let Some(handle) = self
            .pending_state
            .as_mut()
            .and_then(|state| state.ssh_wormhole_timeout_handle.take())
        {
            handle.abort();
        };
    }

    pub fn clear_ssh_block_state(&mut self) {
        if let Some(ref mut pending_state) = self.pending_state.as_mut() {
            pending_state.ssh_block_state = None;
        }
    }

    pub fn set_ssh_block_state(&mut self, ssh_block_state: SshBlockState) {
        let pending_state = self.pending_state.get_or_insert_with(Default::default);
        pending_state.ssh_block_state = Some(ssh_block_state);
    }

    pub fn ssh_block_state(&self) -> Option<&SshBlockState> {
        self.pending_state
            .as_ref()
            .and_then(|state| state.ssh_block_state.as_ref())
    }

    pub fn get_pending_ssh_host(&self) -> Option<String> {
        self.pending_state
            .as_ref()
            .and_then(|state: &WormholeTriggerState| state.pending_wormhole_ssh_host.clone())
    }

    pub fn get_pending_ssh_command(&self) -> Option<String> {
        self.pending_state
            .as_ref()
            .and_then(|state: &WormholeTriggerState| state.pending_command.clone())
    }

    pub fn take_pending_ssh_host(&mut self) -> Option<String> {
        self.pending_state
            .as_mut()
            .and_then(|state: &mut WormholeTriggerState| state.pending_wormhole_ssh_host.take())
    }

    pub fn clear_pending_ssh_host(&mut self) {
        if let Some(ref mut pending_state) = self.pending_state.as_mut() {
            pending_state.pending_wormhole_ssh_host = None;
        }
    }

    pub fn set_pending_ssh_host(&mut self, command: String, ssh_host: Option<String>) {
        let pending_state = self.pending_state.get_or_insert_with(Default::default);
        pending_state.pending_command = Some(command);
        pending_state.pending_wormhole_ssh_host = ssh_host;
    }

    pub fn set_block_id(&mut self, block_id: BlockId) {
        let pending_state = self.pending_state.get_or_insert_with(Default::default);
        pending_state.block_id = Some(block_id);
    }

    pub fn block_id(&self) -> Option<BlockId> {
        self.pending_state
            .as_ref()
            .and_then(|state| state.block_id.clone())
    }

    pub fn timeout_id(&self) -> u8 {
        self.timeout_id
    }

    /// Generates a new timeout ID. This is used to validate that a timeout is still valid.
    /// Call this to get a new timeout ID before starting a new timeout, or to invalidate
    /// an existing timeout.
    pub fn replace_timeout_id(&mut self) -> u8 {
        self.timeout_id = self.timeout_id.wrapping_add(1);
        self.timeout_id
    }

    /// The terminal view should prevent typing
    pub fn should_prevent_input(&self) -> bool {
        let Some(state) = self.ssh_block_state() else {
            return false;
        };
        state.should_prevent_input()
    }

    /// Called once whenever we get a local block completed, as opposed to a remote ssh block
    /// and we have a Wormhole Success block.
    fn on_wormholed_session_complete(
        &mut self,
        state: WormholeTriggerState,
        ctx: &mut ViewContext<TerminalView>,
    ) -> Option<EntityId> {
        self.clear_ssh_block_state();
        ctx.notify();
        let Some(block) = &state.ssh_block_state else {
            return None;
        };
        block.on_wormholed_session_complete(ctx)
    }

    pub fn on_wormhole_start(&mut self, active_session_id: Option<SessionId>) {
        self.session_id = active_session_id;
    }

    /// Called whenever a block is completed, to determine whether a Wormholed session
    /// has been completed.
    pub fn get_completed_wormhole_session_id(
        &mut self,
        active_session_id: Option<SessionId>,
        ctx: &mut ViewContext<TerminalView>,
    ) -> Option<EntityId> {
        if self.session_id.is_none() || active_session_id == self.session_id {
            return None;
        }
        if let Some(state) = self.pending_state.take() {
            return self.on_wormholed_session_complete(state, ctx);
        };
        None
    }
}