pi/interactive/

tree_ui.rs

use super::*;

impl PiApp {
    #[allow(clippy::too_many_lines)]
    pub(super) fn handle_tree_ui_key(&mut self, key: &KeyMsg) -> Option<Cmd> {
        let tree_ui = self.tree_ui.take()?;

        match tree_ui {
            TreeUiState::Selector(mut selector) => {
                match key.key_type {
                    KeyType::Up => selector.move_selection(-1),
                    KeyType::Down => selector.move_selection(1),
                    KeyType::CtrlU => {
                        selector.user_only = !selector.user_only;
                        if let Ok(session_guard) = self.session.try_lock() {
                            selector.rebuild(&session_guard);
                        }
                    }
                    KeyType::CtrlO => {
                        selector.show_all = !selector.show_all;
                        if let Ok(session_guard) = self.session.try_lock() {
                            selector.rebuild(&session_guard);
                        }
                    }
                    KeyType::Esc | KeyType::CtrlC => {
                        self.status_message = Some("Tree navigation cancelled".to_string());
                        self.tree_ui = None;
                        return None;
                    }
                    KeyType::Enter => {
                        if selector.rows.is_empty() {
                            self.tree_ui = None;
                            return None;
                        }

                        let selected = selector.rows[selector.selected].clone();
                        selector.last_selected_id = Some(selected.id.clone());

                        let (new_leaf_id, editor_text) = if let Some(text) = selected.resubmit_text
                        {
                            (selected.parent_id.clone(), Some(text))
                        } else {
                            (Some(selected.id.clone()), None)
                        };

                        // No-op if already at target leaf.
                        if selector.current_leaf_id.as_deref() == new_leaf_id.as_deref() {
                            self.status_message = Some("Already on that branch".to_string());
                            self.tree_ui = None;
                            return None;
                        }

                        let Ok(session_guard) = self.session.try_lock() else {
                            self.status_message = Some("Session busy; try again".to_string());
                            self.tree_ui = None;
                            return None;
                        };

                        let old_leaf_id = session_guard.leaf_id.clone();
                        let (entries_to_summarize, summary_from_id) = collect_tree_branch_entries(
                            &session_guard,
                            old_leaf_id.as_deref(),
                            new_leaf_id.as_deref(),
                        );
                        let session_id = session_guard.header.id.clone();
                        drop(session_guard);

                        let api_key_present = self.agent.try_lock().is_ok_and(|agent_guard| {
                            agent_guard.stream_options().api_key.is_some()
                        });

                        let pending = PendingTreeNavigation {
                            session_id,
                            old_leaf_id,
                            selected_entry_id: selected.id,
                            new_leaf_id,
                            editor_text,
                            entries_to_summarize,
                            summary_from_id,
                            api_key_present,
                        };

                        if pending.entries_to_summarize.is_empty() {
                            // Nothing to summarize; switch immediately.
                            self.start_tree_navigation(pending, TreeSummaryChoice::NoSummary, None);
                            return None;
                        }

                        self.tree_ui = Some(TreeUiState::SummaryPrompt(TreeSummaryPromptState {
                            pending,
                            selected: 0,
                        }));
                        return None;
                    }
                    _ => {}
                }

                self.tree_ui = Some(TreeUiState::Selector(selector));
            }
            TreeUiState::SummaryPrompt(mut prompt) => {
                match key.key_type {
                    KeyType::Up => {
                        if prompt.selected > 0 {
                            prompt.selected -= 1;
                        }
                    }
                    KeyType::Down => {
                        if prompt.selected < TreeSummaryChoice::all().len().saturating_sub(1) {
                            prompt.selected += 1;
                        }
                    }
                    KeyType::Esc | KeyType::CtrlC => {
                        self.status_message = Some("Tree navigation cancelled".to_string());
                        self.tree_ui = None;
                        return None;
                    }
                    KeyType::Enter => {
                        let choice = TreeSummaryChoice::all()[prompt.selected];
                        match choice {
                            TreeSummaryChoice::NoSummary | TreeSummaryChoice::Summarize => {
                                let pending = prompt.pending;
                                self.start_tree_navigation(pending, choice, None);
                                return None;
                            }
                            TreeSummaryChoice::SummarizeWithCustomPrompt => {
                                self.tree_ui =
                                    Some(TreeUiState::CustomPrompt(TreeCustomPromptState {
                                        pending: prompt.pending,
                                        instructions: String::new(),
                                    }));
                                return None;
                            }
                        }
                    }
                    _ => {}
                }
                self.tree_ui = Some(TreeUiState::SummaryPrompt(prompt));
            }
            TreeUiState::CustomPrompt(mut custom) => {
                match key.key_type {
                    KeyType::Esc | KeyType::CtrlC => {
                        self.tree_ui = Some(TreeUiState::SummaryPrompt(TreeSummaryPromptState {
                            pending: custom.pending,
                            selected: 2,
                        }));
                        return None;
                    }
                    KeyType::Backspace => {
                        custom.instructions.pop();
                    }
                    KeyType::Enter => {
                        let pending = custom.pending;
                        let instructions = if custom.instructions.trim().is_empty() {
                            None
                        } else {
                            Some(custom.instructions)
                        };
                        self.start_tree_navigation(
                            pending,
                            TreeSummaryChoice::SummarizeWithCustomPrompt,
                            instructions,
                        );
                        return None;
                    }
                    KeyType::Runes => {
                        for ch in key.runes.iter().copied() {
                            custom.instructions.push(ch);
                        }
                    }
                    _ => {}
                }
                self.tree_ui = Some(TreeUiState::CustomPrompt(custom));
            }
        }

        None
    }

    /// Handle keyboard input when the branch picker overlay is active.
    pub fn handle_branch_picker_key(&mut self, key: &KeyMsg) -> Option<Cmd> {
        let picker = self.branch_picker.as_mut()?;

        match key.key_type {
            KeyType::Up => picker.select_prev(),
            KeyType::Down => picker.select_next(),
            KeyType::Runes if key.runes == ['k'] => picker.select_prev(),
            KeyType::Runes if key.runes == ['j'] => picker.select_next(),
            KeyType::Enter => {
                if let Some(branch) = picker.selected_branch().cloned() {
                    self.branch_picker = None;
                    return self.switch_to_branch_leaf(&branch.leaf_id);
                }
                self.branch_picker = None;
            }
            KeyType::Esc | KeyType::CtrlC => {
                self.branch_picker = None;
                self.status_message = Some("Branch picker cancelled".to_string());
            }
            KeyType::Runes if key.runes == ['q'] => {
                self.branch_picker = None;
            }
            _ => {} // consume all other input while picker is open
        }

        None
    }

    /// Switch the active branch to a different leaf. Reloads the conversation.
    fn switch_to_branch_leaf(&mut self, leaf_id: &str) -> Option<Cmd> {
        let (session_id, old_leaf_id) = self
            .session
            .try_lock()
            .ok()
            .map(|g| (g.header.id.clone(), g.leaf_id.clone()))
            .unwrap_or_default();

        let pending = PendingTreeNavigation {
            session_id,
            old_leaf_id,
            selected_entry_id: leaf_id.to_string(),
            new_leaf_id: Some(leaf_id.to_string()),
            editor_text: None,
            entries_to_summarize: Vec::new(),
            summary_from_id: String::new(),
            api_key_present: false,
        };
        self.start_tree_navigation(pending, TreeSummaryChoice::NoSummary, None);
        None
    }

    /// Open the branch picker if the session has sibling branches.
    pub fn open_branch_picker(&mut self) {
        if self.agent_state != AgentState::Idle {
            self.status_message = Some("Cannot switch branches while processing".to_string());
            return;
        }

        let branches = self
            .session
            .try_lock()
            .ok()
            .and_then(|guard| guard.sibling_branches().map(|(_, b)| b));

        match branches {
            Some(branches) if branches.len() > 1 => {
                self.branch_picker = Some(BranchPickerOverlay::new(branches));
            }
            _ => {
                self.status_message =
                    Some("No branches to pick (use /fork to create one)".to_string());
            }
        }
    }

    /// Cycle to the next or previous sibling branch (Ctrl+Right / Ctrl+Left).
    pub fn cycle_sibling_branch(&mut self, forward: bool) {
        if self.agent_state != AgentState::Idle {
            self.status_message = Some("Cannot switch branches while processing".to_string());
            return;
        }

        let target = self.session.try_lock().ok().and_then(|guard| {
            let (_, branches) = guard.sibling_branches()?;
            if branches.len() <= 1 {
                return None;
            }
            let current_idx = branches.iter().position(|b| b.is_current)?;
            let next_idx = if forward {
                (current_idx + 1) % branches.len()
            } else {
                current_idx.checked_sub(1).unwrap_or(branches.len() - 1)
            };
            Some(branches[next_idx].leaf_id.clone())
        });

        if let Some(leaf_id) = target {
            self.switch_to_branch_leaf(&leaf_id);
        } else {
            self.status_message = Some("No sibling branches (use /fork to create one)".to_string());
        }
    }

    #[allow(clippy::too_many_lines)]
    pub(super) fn start_tree_navigation(
        &mut self,
        pending: PendingTreeNavigation,
        choice: TreeSummaryChoice,
        custom_instructions: Option<String>,
    ) {
        let summary_requested = matches!(
            choice,
            TreeSummaryChoice::Summarize | TreeSummaryChoice::SummarizeWithCustomPrompt
        );

        // Fast path: no summary + no extensions. Keep it synchronous so unit tests can drive it
        // without running the async runtime.
        if !summary_requested && self.extensions.is_none() {
            let Ok(mut session_guard) = self.session.try_lock() else {
                self.status_message = Some("Session busy; try again".to_string());
                return;
            };

            if let Some(target_id) = &pending.new_leaf_id {
                if !session_guard.navigate_to(target_id) {
                    self.status_message = Some(format!("Branch target not found: {target_id}"));
                    return;
                }
            } else {
                session_guard.reset_leaf();
            }

            let (messages, usage) = conversation_from_session(&session_guard);
            let agent_messages = session_guard.to_messages_for_current_path();
            let status_leaf = pending
                .new_leaf_id
                .clone()
                .unwrap_or_else(|| "root".to_string());
            drop(session_guard);

            self.spawn_save_session();

            if let Ok(mut agent_guard) = self.agent.try_lock() {
                agent_guard.replace_messages(agent_messages);
            }

            self.messages = messages;
            self.message_render_cache.clear();
            self.total_usage = usage;
            self.current_response.clear();
            self.current_thinking.clear();
            self.agent_state = AgentState::Idle;
            self.current_tool = None;
            self.abort_handle = None;
            self.status_message = Some(format!("Switched to {status_leaf}"));
            self.scroll_to_bottom();

            if let Some(text) = pending.editor_text {
                self.input.set_value(&text);
            }
            self.input.focus();

            return;
        }

        let event_tx = self.event_tx.clone();
        let session = Arc::clone(&self.session);
        let agent = Arc::clone(&self.agent);
        let extensions = self.extensions.clone();
        let reserve_tokens = self.config.branch_summary_reserve_tokens();
        let runtime_handle = self.runtime_handle.clone();

        let Ok(agent_guard) = self.agent.try_lock() else {
            self.status_message = Some("Agent busy; try again".to_string());
            self.agent_state = AgentState::Idle;
            return;
        };
        let provider = agent_guard.provider();
        let key_opt = agent_guard.stream_options().api_key.clone();

        self.tree_ui = None;
        self.agent_state = AgentState::Processing;
        self.status_message = Some("Switching branches...".to_string());

        runtime_handle.spawn(async move {
            let cx = Cx::for_request();

            let from_id_for_event = pending
                .old_leaf_id
                .clone()
                .unwrap_or_else(|| "root".to_string());
            let to_id_for_event = pending
                .new_leaf_id
                .clone()
                .unwrap_or_else(|| "root".to_string());

            if let Some(manager) = extensions.clone() {
                let cancelled = manager
                    .dispatch_cancellable_event(
                        ExtensionEventName::SessionBeforeSwitch,
                        Some(json!({
                            "fromId": from_id_for_event.clone(),
                            "toId": to_id_for_event.clone(),
                            "sessionId": pending.session_id.clone(),
                        })),
                        EXTENSION_EVENT_TIMEOUT_MS,
                    )
                    .await
                    .unwrap_or(false);
                if cancelled {
                    let _ = event_tx.try_send(PiMsg::System(
                        "Session switch cancelled by extension".to_string(),
                    ));
                    return;
                }
            }

            let summary_skipped =
                summary_requested && key_opt.is_none() && !pending.entries_to_summarize.is_empty();
            let summary_text = if !summary_requested || pending.entries_to_summarize.is_empty() {
                None
            } else if let Some(key) = key_opt.as_deref() {
                match crate::compaction::summarize_entries(
                    &pending.entries_to_summarize,
                    provider,
                    key,
                    reserve_tokens,
                    custom_instructions.as_deref(),
                )
                .await
                {
                    Ok(summary) => summary,
                    Err(err) => {
                        let _ = event_tx
                            .try_send(PiMsg::AgentError(format!("Branch summary failed: {err}")));
                        return;
                    }
                }
            } else {
                None
            };

            let messages_for_agent = {
                let mut guard = match session.lock(&cx).await {
                    Ok(guard) => guard,
                    Err(err) => {
                        let _ = event_tx
                            .try_send(PiMsg::AgentError(format!("Failed to lock session: {err}")));
                        return;
                    }
                };

                if let Some(target_id) = &pending.new_leaf_id {
                    if !guard.navigate_to(target_id) {
                        let _ = event_tx.try_send(PiMsg::AgentError(format!(
                            "Branch target not found: {target_id}"
                        )));
                        return;
                    }
                } else {
                    guard.reset_leaf();
                }

                if let Some(summary_text) = summary_text {
                    guard.append_branch_summary(
                        pending.summary_from_id.clone(),
                        summary_text,
                        None,
                        None,
                    );
                }

                let _ = guard.save().await;
                guard.to_messages_for_current_path()
            };

            {
                let mut agent_guard = match agent.lock(&cx).await {
                    Ok(guard) => guard,
                    Err(err) => {
                        let _ = event_tx
                            .try_send(PiMsg::AgentError(format!("Failed to lock agent: {err}")));
                        return;
                    }
                };
                agent_guard.replace_messages(messages_for_agent);
            }

            let (messages, usage) = {
                let guard = match session.lock(&cx).await {
                    Ok(guard) => guard,
                    Err(err) => {
                        let _ = event_tx
                            .try_send(PiMsg::AgentError(format!("Failed to lock session: {err}")));
                        return;
                    }
                };
                conversation_from_session(&guard)
            };

            let status = if summary_skipped {
                Some(format!(
                    "Switched to {to_id_for_event} (no summary: missing API key)"
                ))
            } else {
                Some(format!("Switched to {to_id_for_event}"))
            };

            let _ = event_tx.try_send(PiMsg::ConversationReset {
                messages,
                usage,
                status,
            });

            if let Some(text) = pending.editor_text {
                let _ = event_tx.try_send(PiMsg::SetEditorText(text));
            }

            if let Some(manager) = extensions {
                let _ = manager
                    .dispatch_event(
                        ExtensionEventName::SessionSwitch,
                        Some(json!({
                            "fromId": from_id_for_event,
                            "toId": to_id_for_event,
                            "sessionId": pending.session_id,
                        })),
                    )
                    .await;
            }
        });
    }
}