oxi-cli 0.4.2

//! Session tree navigation for branched sessions
//!
//! Handles navigation between different points in the session tree,
//! including the creation of branch summaries when leaving a path.

use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use uuid::Uuid;

// ============================================================================
// Types
// ============================================================================

/// Extended session entry types for tree navigation
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum SessionEntryType {
    /// A regular message (user/assistant/system/tool)
    Message(MessageEntry),
    /// Summary of a branch when navigating away
    BranchSummary(BranchSummaryEntry),
    /// Summary after context compaction
    Compaction(CompactionEntry),
    /// A label/marker on an entry
    Label(LabelEntry),
    /// User-defined metadata for the session
    SessionInfo(SessionInfoEntry),
    /// Extension-specific data (not sent to LLM)
    Custom(CustomEntry),
    /// Custom message from extension (sent to LLM)
    CustomMessage(CustomMessageEntry),
}

/// Message entry content
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageEntry {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub timestamp: i64,
    pub role: MessageRole,
    pub content: String,
}

/// Message role types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum MessageRole {
    User,
    Assistant,
    System,
    Tool,
    ToolResult,
    Custom,
}

impl MessageRole {
    /// Check if this is a user message
    pub fn is_user(&self) -> bool {
        matches!(self, MessageRole::User)
    }

    /// Check if this is an assistant message
    pub fn is_assistant(&self) -> bool {
        matches!(self, MessageRole::Assistant)
    }
}

/// Branch summary entry - captures context when navigating away from a branch
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BranchSummaryEntry {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub timestamp: i64,
    /// ID of the entry this summary is attached to
    pub from_id: Uuid,
    /// The summary text
    pub summary: String,
    /// Optional details (e.g., file lists)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub details: Option<BranchSummaryDetails>,
    /// True if generated by an extension
    #[serde(skip_serializing_if = "Option::is_none")]
    pub from_hook: Option<bool>,
}

/// Details stored in BranchSummaryEntry for file tracking
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BranchSummaryDetails {
    pub read_files: Vec<String>,
    pub modified_files: Vec<String>,
}

/// Compaction entry - summary after context window compaction
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompactionEntry {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub timestamp: i64,
    pub summary: String,
    /// ID of first entry kept after compaction
    pub first_kept_entry_id: Uuid,
    /// Estimated tokens before compaction
    pub tokens_before: usize,
    /// Extension-specific data
    #[serde(skip_serializing_if = "Option::is_none")]
    pub details: Option<serde_json::Value>,
    /// True if generated by an extension
    #[serde(skip_serializing_if = "Option::is_none")]
    pub from_hook: Option<bool>,
}

/// Label entry - user-defined bookmark/marker on entries
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LabelEntry {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub timestamp: i64,
    pub target_id: Uuid,
    pub label: Option<String>,
}

/// Session info entry - metadata like display name
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionInfoEntry {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub timestamp: i64,
    pub name: Option<String>,
}

/// Custom entry for extensions to store extension-specific data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CustomEntry {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub timestamp: i64,
    pub custom_type: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<serde_json::Value>,
}

/// Custom message entry for extensions to inject messages into LLM context
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CustomMessageEntry {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub timestamp: i64,
    pub custom_type: String,
    pub content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub details: Option<serde_json::Value>,
    /// Controls TUI rendering
    pub display: bool,
}

/// Result of collecting entries for branch summarization
#[derive(Debug, Clone)]
pub struct CollectEntriesResult {
    /// Entries to summarize, in chronological order
    pub entries: Vec<SessionEntryType>,
    /// Common ancestor between old and new position
    pub common_ancestor_id: Option<Uuid>,
}

/// Navigation target options
#[derive(Debug, Clone)]
pub struct NavigationOptions {
    /// If true, generate a summary of the branch being left
    pub summarize: bool,
    /// Custom instructions for summarization
    pub custom_instructions: Option<String>,
    /// If true, custom_instructions replaces the default prompt instead of being appended
    pub replace_instructions: bool,
    /// Label to attach to the branch summary entry
    pub label: Option<String>,
}

impl Default for NavigationOptions {
    fn default() -> Self {
        Self {
            summarize: false,
            custom_instructions: None,
            replace_instructions: false,
            label: None,
        }
    }
}

/// Result of tree navigation
#[derive(Debug, Clone)]
pub struct NavigationResult {
    /// Text content for the editor (user message text when navigating to user message)
    pub editor_text: Option<String>,
    /// True if navigation was cancelled
    pub cancelled: bool,
    /// True if summarization was aborted
    pub aborted: bool,
    /// ID of the summary entry created, if any
    pub summary_entry_id: Option<Uuid>,
}

/// Preparation data for tree navigation (passed to extension hooks)
#[derive(Debug, Clone)]
pub struct TreePreparation {
    pub target_id: Uuid,
    pub old_leaf_id: Option<Uuid>,
    pub common_ancestor_id: Option<Uuid>,
    pub entries_to_summarize: Vec<SessionEntryType>,
    pub user_wants_summary: bool,
    pub custom_instructions: Option<String>,
    pub replace_instructions: bool,
    pub label: Option<String>,
}

/// Trait for LLM-based summarization
/// Implement this to provide actual summarization, or use a mock for testing
pub trait Summarizer: Send + Sync {
    /// Generate a summary of the given entries
    fn summarize(
        &self,
        entries: &[SessionEntryType],
        custom_instructions: Option<&str>,
        replace_instructions: bool,
    ) -> impl std::future::Future<Output = Result<BranchSummaryResult, SummarizationError>>
           + Send
           + 'static;
}

/// Result of branch summarization
#[derive(Debug, Clone)]
pub struct BranchSummaryResult {
    /// The generated summary text
    pub summary: Option<String>,
    /// Files that were read during this branch
    pub read_files: Vec<String>,
    /// Files that were modified during this branch
    pub modified_files: Vec<String>,
    /// True if the operation was aborted
    pub aborted: bool,
    /// Error message if summarization failed
    pub error: Option<String>,
}

/// Error during summarization
#[derive(Debug, Clone)]
pub enum SummarizationError {
    NoModel,
    Aborted,
    Failed(String),
}

/// Extension hook result for session_before_tree event
#[derive(Debug, Clone)]
pub struct BeforeTreeHookResult {
    /// If true, cancel the navigation
    pub cancel: bool,
    /// Summary from extension (if user wants summary)
    pub summary: Option<ExtensionSummary>,
    /// Override custom instructions
    pub custom_instructions: Option<String>,
    /// Override replace instructions flag
    pub replace_instructions: Option<bool>,
    /// Override label
    pub label: Option<String>,
}

/// Summary generated by an extension
#[derive(Debug, Clone)]
pub struct ExtensionSummary {
    pub summary: String,
    pub details: Option<serde_json::Value>,
}

// ============================================================================
// Tree Navigation Implementation
// ============================================================================

/// Session navigator for tree-based session management
pub struct SessionNavigator {
    /// All entries indexed by ID
    entries_by_id: HashMap<Uuid, SessionEntryType>,
    /// Labels by entry ID
    labels_by_id: HashMap<Uuid, String>,
    /// Label timestamps by entry ID
    label_timestamps_by_id: HashMap<Uuid, i64>,
    /// Current leaf ID (where new entries are appended)
    leaf_id: Option<Uuid>,
}

impl SessionNavigator {
    /// Create a new session navigator
    pub fn new() -> Self {
        Self {
            entries_by_id: HashMap::new(),
            labels_by_id: HashMap::new(),
            label_timestamps_by_id: HashMap::new(),
            leaf_id: None,
        }
    }

    /// Create from existing entries
    pub fn from_entries(entries: Vec<SessionEntryType>, leaf_id: Option<Uuid>) -> Self {
        let mut entries_by_id = HashMap::new();
        for entry in &entries {
            let id = Self::entry_id(entry);
            entries_by_id.insert(id, entry.clone());
        }

        Self {
            entries_by_id,
            labels_by_id: HashMap::new(),
            label_timestamps_by_id: HashMap::new(),
            leaf_id,
        }
    }

    /// Get the current leaf ID
    pub fn get_leaf_id(&self) -> Option<Uuid> {
        self.leaf_id
    }

    /// Get an entry by ID
    pub fn get_entry(&self, id: Uuid) -> Option<&SessionEntryType> {
        self.entries_by_id.get(&id)
    }

    /// Get the label for an entry, if any
    pub fn get_label(&self, id: Uuid) -> Option<&str> {
        self.labels_by_id.get(&id).map(|s| s.as_str())
    }

    /// Get all entries
    pub fn get_entries(&self) -> Vec<&SessionEntryType> {
        self.entries_by_id.values().collect()
    }

    /// Get the branch path from root to a given entry
    pub fn get_branch(&self, from_id: Option<Uuid>) -> Vec<&SessionEntryType> {
        let mut path = Vec::new();
        let start_id = from_id.or(self.leaf_id);

        let mut current_id = start_id;
        while let Some(id) = current_id {
            if let Some(entry) = self.entries_by_id.get(&id) {
                path.insert(0, entry);
                current_id = Self::entry_parent_id(entry);
            } else {
                break;
            }
        }

        path
    }

    /// Get all direct children of an entry
    pub fn get_children(&self, parent_id: Uuid) -> Vec<&SessionEntryType> {
        self.entries_by_id
            .values()
            .filter(|entry| Self::entry_parent_id(entry) == Some(parent_id))
            .collect()
    }

    /// Collect entries that should be summarized when navigating from one position to another.
    ///
    /// Walks from oldLeafId back to the common ancestor with targetId, collecting entries
    /// along the way.
    pub fn collect_entries_for_branch_summary(
        &self,
        old_leaf_id: Option<Uuid>,
        target_id: Uuid,
    ) -> CollectEntriesResult {
        // If no old position, nothing to summarize
        if old_leaf_id.is_none() {
            return CollectEntriesResult {
                entries: Vec::new(),
                common_ancestor_id: None,
            };
        }

        let old_leaf_id = old_leaf_id.unwrap();

        // Build sets of IDs on both paths
        let old_path_ids: HashSet<Uuid> = self
            .get_branch(Some(old_leaf_id))
            .iter()
            .map(|e| Self::entry_id(e))
            .collect();

        let target_path = self.get_branch(Some(target_id));

        // Find common ancestor (deepest node that's on both paths)
        // target_path is root-first, so iterate backwards
        let mut common_ancestor_id: Option<Uuid> = None;
        for entry in target_path.iter().rev() {
            let id = Self::entry_id(entry);
            if old_path_ids.contains(&id) {
                common_ancestor_id = Some(id);
                break;
            }
        }

        // Collect entries from old leaf back to common ancestor
        let mut entries: Vec<&SessionEntryType> = Vec::new();
        let mut current_id: Option<Uuid> = Some(old_leaf_id);

        while let Some(id) = current_id {
            if current_id == common_ancestor_id {
                break;
            }
            if let Some(entry) = self.entries_by_id.get(&id) {
                entries.push(entry);
                current_id = Self::entry_parent_id(entry);
            } else {
                break;
            }
        }

        // Reverse to get chronological order
        entries.reverse();

        CollectEntriesResult {
            entries: entries.into_iter().cloned().collect(),
            common_ancestor_id,
        }
    }

    /// Navigate to a target entry in the session tree.
    ///
    /// This handles:
    /// - No-op if already at target
    /// - Collection of entries to summarize
    /// - Extension hooks (session_before_tree, session_tree)
    /// - Branch summarization via the provided summarizer
    /// - Leaf switching based on target type
    /// - Label attachment
    pub fn navigate_tree<N: Summarizer + ?Sized>(
        &mut self,
        target_id: Uuid,
        options: NavigationOptions,
        summarizer: Option<&N>,
        extension_hook: Option<&dyn Fn(TreePreparation) -> BeforeTreeHookResult>,
    ) -> NavigationResult {
        let old_leaf_id = self.leaf_id;

        // No-op if already at target
        if Some(target_id) == old_leaf_id {
            return NavigationResult {
                editor_text: None,
                cancelled: false,
                aborted: false,
                summary_entry_id: None,
            };
        }

        // Validate target exists
        let target_entry = match self.entries_by_id.get(&target_id) {
            Some(e) => e,
            None => {
                return NavigationResult {
                    editor_text: None,
                    cancelled: true,
                    aborted: false,
                    summary_entry_id: None,
                };
            }
        };

        // Collect entries to summarize (from old leaf to common ancestor)
        let collection = self.collect_entries_for_branch_summary(old_leaf_id, target_id);

        // Prepare mutable options that can be overridden by extensions
        let mut custom_instructions = options.custom_instructions.clone();
        let mut replace_instructions = options.replace_instructions;
        let mut label = options.label.clone();

        // Prepare tree preparation for extension hooks
        let preparation = TreePreparation {
            target_id,
            old_leaf_id,
            common_ancestor_id: collection.common_ancestor_id,
            entries_to_summarize: collection.entries.clone(),
            user_wants_summary: options.summarize,
            custom_instructions: custom_instructions.clone(),
            replace_instructions,
            label: label.clone(),
        };

        // Run extension hook (session_before_tree)
        let mut extension_summary: Option<ExtensionSummary> = None;
        let mut from_extension = false;

        if let Some(hook) = extension_hook {
            let result = hook(preparation);
            if result.cancel {
                return NavigationResult {
                    editor_text: None,
                    cancelled: true,
                    aborted: false,
                    summary_entry_id: None,
                };
            }

            if let Some(ext_sum) = result.summary {
                extension_summary = Some(ext_sum);
                from_extension = true;
            }

            // Allow extensions to override instructions and label
            if let Some(ci) = result.custom_instructions {
                custom_instructions = Some(ci);
            }
            if let Some(ri) = result.replace_instructions {
                replace_instructions = ri;
            }
            if let Some(l) = result.label {
                label = Some(l);
            }
        }

        // Run default summarizer if needed
        let mut summary_text: Option<String> = None;
        let mut summary_details: Option<BranchSummaryDetails> = None;

        if options.summarize && !collection.entries.is_empty() && extension_summary.is_none() {
            if let Some(summarizer) = summarizer {
                let rt = summarizer.summarize(
                    &collection.entries,
                    custom_instructions.as_deref(),
                    replace_instructions,
                );
                let runtime = tokio::runtime::Handle::current();
                let result = runtime.block_on(rt);

                match result {
                    Ok(summary_result) => {
                        if summary_result.aborted {
                            return NavigationResult {
                                editor_text: None,
                                cancelled: true,
                                aborted: true,
                                summary_entry_id: None,
                            };
                        }
                        if let Some(err) = summary_result.error {
                            // Log error but continue without summary
                            tracing::warn!("Summarization failed: {}", err);
                        }
                        summary_text = summary_result.summary;
                        if !summary_result.read_files.is_empty()
                            || !summary_result.modified_files.is_empty()
                        {
                            summary_details = Some(BranchSummaryDetails {
                                read_files: summary_result.read_files,
                                modified_files: summary_result.modified_files,
                            });
                        }
                    }
                    Err(e) => {
                        tracing::warn!("Summarization error: {:?}", e);
                    }
                }
            }
        } else if let Some(ext_sum) = extension_summary {
            summary_text = Some(ext_sum.summary);
            summary_details = ext_sum.details.and_then(|d| {
                serde_json::from_value(d).ok()
            });
        }

        // Determine the new leaf position based on target type
        let (new_leaf_id, editor_text) = Self::determine_leaf_and_editor(target_entry);

        // Switch leaf (with or without summary)
        let summary_entry_id = if let Some(text) = summary_text {
            // Create summary at target position (can be null for root)
            let summary_id = self.branch_with_summary(new_leaf_id, text, summary_details, from_extension);

            // Attach label to the summary entry
            if let Some(l) = &label {
                self.append_label_change(summary_id, Some(l.clone()));
            }

            Some(summary_id)
        } else if new_leaf_id.is_none() {
            // No summary, navigating to root - reset leaf
            self.reset_leaf();
            None
        } else {
            // No summary, navigating to non-root
            self.branch(new_leaf_id.unwrap());
            None
        };

        // Attach label to target entry when not summarizing
        let has_label = label.is_some();
        let has_summary = summary_text.is_some();
        if has_label && !has_summary {
            self.append_label_change(target_id, label);
        }

        NavigationResult {
            editor_text,
            cancelled: false,
            aborted: false,
            summary_entry_id,
        }
    }

    /// Determine the new leaf position and editor text based on target entry type.
    ///
    /// - User message: leaf = parent (null if root), text goes to editor
    /// - Custom message: leaf = parent (null if root), text goes to editor
    /// - Other entries: leaf = selected node
    fn determine_leaf_and_editor(entry: &SessionEntryType) -> (Option<Uuid>, Option<String>) {
        match entry {
            SessionEntryType::Message(msg) if msg.role.is_user() => {
                let editor_text = if msg.content.is_empty() {
                    None
                } else {
                    Some(msg.content.clone())
                };
                (msg.parent_id, editor_text)
            }
            SessionEntryType::CustomMessage(custom) => {
                let editor_text = if custom.content.is_empty() {
                    None
                } else {
                    Some(custom.content.clone())
                };
                (custom.parent_id, editor_text)
            }
            _ => {
                let leaf_id = Self::entry_id(entry);
                (Some(leaf_id), None)
            }
        }
    }

    /// Switch to a different entry (start a new branch).
    pub fn branch(&mut self, branch_from_id: Uuid) {
        if !self.entries_by_id.contains_key(&branch_from_id) {
            tracing::warn!("Entry {} not found for branching", branch_from_id);
            return;
        }
        self.leaf_id = Some(branch_from_id);
    }

    /// Reset the leaf pointer to null (before any entries).
    pub fn reset_leaf(&mut self) {
        self.leaf_id = None;
    }

    /// Start a new branch with a summary of the abandoned path.
    /// Returns the ID of the new summary entry.
    pub fn branch_with_summary(
        &mut self,
        branch_from_id: Option<Uuid>,
        summary: String,
        details: Option<BranchSummaryDetails>,
        from_hook: bool,
    ) -> Uuid {
        // Validate branch_from_id exists if provided
        if branch_from_id.is_some()
            && !self.entries_by_id.contains_key(&branch_from_id.unwrap())
        {
            tracing::warn!(
                "Entry {:?} not found for branch_with_summary",
                branch_from_id
            );
        }

        self.leaf_id = branch_from_id;

        let summary_id = Uuid::new_v4();
        let entry = SessionEntryType::BranchSummary(BranchSummaryEntry {
            id: summary_id,
            parent_id: branch_from_id,
            timestamp: Utc::now().timestamp_millis(),
            from_id: branch_from_id.unwrap_or_else(|| Uuid::nil()),
            summary,
            details,
            from_hook: Some(from_hook),
        });

        self.entries_by_id.insert(summary_id, entry.clone());
        summary_id
    }

    /// Set or clear a label on an entry.
    /// Returns the ID of the label entry created.
    pub fn append_label_change(&mut self, target_id: Uuid, label: Option<String>) -> Uuid {
        if !self.entries_by_id.contains_key(&target_id) {
            tracing::warn!("Entry {} not found for label change", target_id);
            return Uuid::nil();
        }

        let label_id = Uuid::new_v4();
        let entry = SessionEntryType::Label(LabelEntry {
            id: label_id,
            parent_id: self.leaf_id,
            timestamp: Utc::now().timestamp_millis(),
            target_id,
            label: label.clone(),
        });

        self.entries_by_id.insert(label_id, entry);

        if let Some(l) = label {
            self.labels_by_id.insert(target_id, l);
            self.label_timestamps_by_id
                .insert(target_id, Utc::now().timestamp_millis());
        } else {
            self.labels_by_id.remove(&target_id);
            self.label_timestamps_by_id.remove(&target_id);
        }

        label_id
    }

    /// Add an entry to the session
    pub fn add_entry(&mut self, entry: SessionEntryType) {
        let id = Self::entry_id(&entry);
        self.entries_by_id.insert(id, entry);
    }

    /// Get the label timestamp for an entry, if any
    pub fn get_label_timestamp(&self, id: Uuid) -> Option<i64> {
        self.label_timestamps_by_id.get(&id).copied()
    }

    // Helper methods for accessing entry fields generically

    fn entry_id(entry: &SessionEntryType) -> Uuid {
        match entry {
            SessionEntryType::Message(e) => e.id,
            SessionEntryType::BranchSummary(e) => e.id,
            SessionEntryType::Compaction(e) => e.id,
            SessionEntryType::Label(e) => e.id,
            SessionEntryType::SessionInfo(e) => e.id,
            SessionEntryType::Custom(e) => e.id,
            SessionEntryType::CustomMessage(e) => e.id,
        }
    }

    fn entry_parent_id(entry: &SessionEntryType) -> Option<Uuid> {
        match entry {
            SessionEntryType::Message(e) => e.parent_id,
            SessionEntryType::BranchSummary(e) => e.parent_id,
            SessionEntryType::Compaction(e) => e.parent_id,
            SessionEntryType::Label(e) => e.parent_id,
            SessionEntryType::SessionInfo(e) => e.parent_id,
            SessionEntryType::Custom(e) => e.parent_id,
            SessionEntryType::CustomMessage(e) => e.parent_id,
        }
    }
}

impl Default for SessionNavigator {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Utility Functions
// ============================================================================

/// Extract text content from user message
pub fn extract_user_message_text(content: &str) -> String {
    content.to_string()
}

/// Extract text content from custom message
pub fn extract_custom_message_text(content: &str) -> String {
    content.to_string()
}

/// Check if an entry is a user message
pub fn is_user_message(entry: &SessionEntryType) -> bool {
    matches!(
        entry,
        SessionEntryType::Message(msg) if msg.role.is_user()
    )
}

/// Check if an entry is a custom message
pub fn is_custom_message(entry: &SessionEntryType) -> bool {
    matches!(entry, SessionEntryType::CustomMessage(_))
}

/// Check if an entry is an assistant message
pub fn is_assistant_message(entry: &SessionEntryType) -> bool {
    matches!(
        entry,
        SessionEntryType::Message(msg) if msg.role.is_assistant()
    )
}

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

    fn create_message(id: Uuid, parent_id: Option<Uuid>, role: MessageRole, content: &str) -> SessionEntryType {
        SessionEntryType::Message(MessageEntry {
            id,
            parent_id,
            timestamp: 0,
            role,
            content: content.to_string(),
        })
    }

    #[test]
    fn test_navigate_to_user_message() {
        let mut nav = SessionNavigator::new();

        // Create a simple tree: root -> user -> assistant
        let root_id = Uuid::new_v4();
        let user_id = Uuid::new_v4();
        let assistant_id = Uuid::new_v4();

        nav.add_entry(create_message(root_id, None, MessageRole::User, "Hello"));
        nav.add_entry(create_message(user_id, Some(root_id), MessageRole::User, "How are you?"));
        nav.add_entry(create_message(assistant_id, Some(user_id), MessageRole::Assistant, "I'm fine"));

        nav.branch(assistant_id);

        // Navigate back to user message
        let result = nav.navigate_tree(user_id, NavigationOptions::default(), None::<&dyn Summarizer>, None);

        assert!(!result.cancelled);
        assert!(!result.aborted);
        assert_eq!(result.editor_text, Some("How are you?".to_string()));
        assert_eq!(nav.get_leaf_id(), Some(user_id));
    }

    #[test]
    fn test_navigate_to_assistant_message() {
        let mut nav = SessionNavigator::new();

        let root_id = Uuid::new_v4();
        let user_id = Uuid::new_v4();
        let assistant_id = Uuid::new_v4();

        nav.add_entry(create_message(root_id, None, MessageRole::User, "Hello"));
        nav.add_entry(create_message(user_id, Some(root_id), MessageRole::User, "How are you?"));
        nav.add_entry(create_message(assistant_id, Some(user_id), MessageRole::Assistant, "I'm fine"));

        nav.branch(assistant_id);

        // Navigate to assistant message (should stay there)
        let result = nav.navigate_tree(assistant_id, NavigationOptions::default(), None::<&dyn Summarizer>, None);

        assert!(!result.cancelled);
        assert_eq!(nav.get_leaf_id(), Some(assistant_id));
    }

    #[test]
    fn test_noop_navigation() {
        let mut nav = SessionNavigator::new();

        let entry_id = Uuid::new_v4();
        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));
        nav.branch(entry_id);

        // Navigate to same position (no-op)
        let result = nav.navigate_tree(entry_id, NavigationOptions::default(), None::<&dyn Summarizer>, None);

        assert!(!result.cancelled);
        assert_eq!(result.editor_text, None);
    }

    #[test]
    fn test_collect_entries_for_branch_summary() {
        let nav = SessionNavigator::new();

        let root_id = Uuid::new_v4();
        let user_id = Uuid::new_v4();
        let assistant_id = Uuid::new_v4();
        let branch_user_id = Uuid::new_v4();
        let branch_assistant_id = Uuid::new_v4();

        // Main branch: root -> user -> assistant
        let mut entries = vec![
            create_message(root_id, None, MessageRole::User, "Root"),
            create_message(user_id, Some(root_id), MessageRole::User, "User"),
            create_message(assistant_id, Some(user_id), MessageRole::Assistant, "Assistant"),
        ];

        // Branch from user: user -> branch_user -> branch_assistant
        entries.push(create_message(branch_user_id, Some(user_id), MessageRole::User, "Branch User"));
        entries.push(create_message(branch_assistant_id, Some(branch_user_id), MessageRole::Assistant, "Branch Assistant"));

        let nav = SessionNavigator::from_entries(entries, Some(branch_assistant_id));

        // Collect entries from branch_assistant back to root
        let result = nav.collect_entries_for_branch_summary(Some(branch_assistant_id), root_id);

        assert_eq!(result.common_ancestor_id, Some(root_id));
        // Should collect branch_assistant and branch_user
        assert_eq!(result.entries.len(), 2);
    }

    #[test]
    fn test_label_attachment() {
        let mut nav = SessionNavigator::new();

        let entry_id = Uuid::new_v4();
        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));

        // Add label
        let label_id = nav.append_label_change(entry_id, Some("Important".to_string()));

        assert_eq!(nav.get_label(entry_id), Some("Important"));
        assert!(nav.get_entry(label_id).is_some());

        // Remove label
        nav.append_label_change(entry_id, None);

        assert_eq!(nav.get_label(entry_id), None);
    }

    #[test]
    fn test_branch_with_summary() {
        let mut nav = SessionNavigator::new();

        let entry_id = Uuid::new_v4();
        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));

        nav.branch(entry_id);

        // Create a branch with summary
        let summary_id = nav.branch_with_summary(
            Some(entry_id),
            "This is a summary".to_string(),
            None,
            false,
        );

        assert!(nav.get_entry(summary_id).is_some());
        assert_eq!(nav.get_leaf_id(), Some(entry_id));

        // Verify it's a branch summary entry
        match nav.get_entry(summary_id) {
            Some(SessionEntryType::BranchSummary(e)) => {
                assert_eq!(e.summary, "This is a summary");
            }
            _ => panic!("Expected branch summary entry"),
        }
    }
}