oxi-store 0.16.0

//! 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 {
    /// Unique identifier for this message.
    pub id: Uuid,
    /// Parent entry ID in the session tree.
    pub parent_id: Option<Uuid>,
    /// Unix timestamp in milliseconds.
    pub timestamp: i64,
    /// Role of the message sender.
    pub role: MessageRole,
    /// Text content of the message.
    pub content: String,
}

/// Message role types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum MessageRole {
    /// User-authored message.
    User,
    /// Assistant-authored message.
    Assistant,
    /// System prompt message.
    System,
    /// Tool call message.
    Tool,
    /// Tool result message.
    ToolResult,
    /// Custom role message.
    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 {
    /// Unique identifier for this summary entry.
    pub id: Uuid,
    /// Parent entry ID in the session tree.
    pub parent_id: Option<Uuid>,
    /// Unix timestamp in milliseconds.
    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 {
    /// Files that were read during the branch.
    pub read_files: Vec<String>,
    /// Files that were modified during the branch.
    pub modified_files: Vec<String>,
}

/// Compaction entry - summary after context window compaction
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompactionEntry {
    /// Unique identifier for this compaction entry.
    pub id: Uuid,
    /// Parent entry ID in the session tree.
    pub parent_id: Option<Uuid>,
    /// Unix timestamp in milliseconds.
    pub timestamp: i64,
    /// Compacted summary text.
    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 {
    /// Unique identifier for this label entry.
    pub id: Uuid,
    /// Parent entry ID in the session tree.
    pub parent_id: Option<Uuid>,
    /// Unix timestamp in milliseconds.
    pub timestamp: i64,
    /// Entry ID this label is attached to.
    pub target_id: Uuid,
    /// Optional label text.
    pub label: Option<String>,
}

/// Session info entry - metadata like display name
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionInfoEntry {
    /// Unique identifier for this session info entry.
    pub id: Uuid,
    /// Parent entry ID in the session tree.
    pub parent_id: Option<Uuid>,
    /// Unix timestamp in milliseconds.
    pub timestamp: i64,
    /// Optional display name for the session.
    pub name: Option<String>,
}

/// Custom entry for extensions to store extension-specific data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CustomEntry {
    /// Unique identifier for this custom entry.
    pub id: Uuid,
    /// Parent entry ID in the session tree.
    pub parent_id: Option<Uuid>,
    /// Unix timestamp in milliseconds.
    pub timestamp: i64,
    /// Extension-defined type discriminator.
    pub custom_type: String,
    /// Optional extension-specific data payload.
    #[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 {
    /// Unique identifier for this custom message entry.
    pub id: Uuid,
    /// Parent entry ID in the session tree.
    pub parent_id: Option<Uuid>,
    /// Unix timestamp in milliseconds.
    pub timestamp: i64,
    /// Extension-defined type discriminator.
    pub custom_type: String,
    /// Message content injected into LLM context.
    pub content: String,
    /// Optional extension-specific details.
    #[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, Default)]
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>,
}

/// 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 {
    /// The navigation target entry ID.
    pub target_id: Uuid,
    /// The leaf ID before navigation.
    pub old_leaf_id: Option<Uuid>,
    /// Common ancestor between old and new positions.
    pub common_ancestor_id: Option<Uuid>,
    /// Entries collected for summarization.
    pub entries_to_summarize: Vec<SessionEntryType>,
    /// Whether the user requested a summary.
    pub user_wants_summary: bool,
    /// Custom instructions for summarization.
    pub custom_instructions: Option<String>,
    /// Whether custom instructions replace the default prompt.
    pub replace_instructions: bool,
    /// Label to attach to the summary.
    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 {
    /// No model available for summarization.
    NoModel,
    /// Summarization was aborted by the user.
    Aborted,
    /// Summarization failed with an error message.
    Failed(String),
}

impl std::fmt::Display for SummarizationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SummarizationError::NoModel => write!(f, "No model available for summarization"),
            SummarizationError::Aborted => write!(f, "Summarization was aborted"),
            SummarizationError::Failed(msg) => write!(f, "Summarization failed: {}", msg),
        }
    }
}

/// 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 {
    /// Summary text generated by the extension.
    pub summary: String,
    /// Optional extension-specific details.
    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.
    pub fn collect_entries_for_branch_summary(
        &self,
        old_leaf_id: Option<Uuid>,
        target_id: Uuid,
    ) -> CollectEntriesResult {
        let old_leaf_id = match old_leaf_id {
            Some(id) => id,
            None => {
                return CollectEntriesResult {
                    entries: Vec::new(),
                    common_ancestor_id: None,
                };
            }
        };

        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));

        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;
            }
        }

        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;
            }
        }

        entries.reverse();

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

    /// Navigate to a target entry in the session tree.
    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;

        if Some(target_id) == old_leaf_id {
            return NavigationResult {
                editor_text: None,
                cancelled: false,
                aborted: false,
                summary_entry_id: None,
            };
        }

        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,
                };
            }
        };

        let collection = self.collect_entries_for_branch_summary(old_leaf_id, target_id);

        let mut custom_instructions = options.custom_instructions.clone();
        let mut replace_instructions = options.replace_instructions;
        let mut label = options.label.clone();

        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(),
        };

        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;
            }

            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);
            }
        }

        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 {
                // Use tokio::task::block_in_place to safely call block_on from within
                // a multi-threaded tokio runtime. Falls back to creating a new runtime
                // if not in a tokio context.
                let entries_clone: Vec<SessionEntryType> = collection.entries.clone();
                let custom_clone = custom_instructions.clone();
                let result = tokio::task::block_in_place(|| {
                    tokio::runtime::Handle::current().block_on(summarizer.summarize(
                        &entries_clone,
                        custom_clone.as_deref(),
                        replace_instructions,
                    ))
                });

                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 {
                            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());
        }

        let (new_leaf_id, editor_text) = Self::determine_leaf_and_editor(target_entry);

        let has_summary = summary_text.is_some();
        let summary_entry_id = if let Some(text) = summary_text {
            let summary_id =
                self.branch_with_summary(new_leaf_id, text, summary_details, from_extension);

            if let Some(l) = &label {
                self.append_label_change(summary_id, Some(l.clone()));
            }

            Some(summary_id)
        } else if new_leaf_id.is_none() {
            self.reset_leaf();
            None
        } else {
            if let Some(id) = new_leaf_id {
                self.branch(id);
            } else {
                self.reset_leaf();
            }
            None
        };

        let has_label = label.is_some();
        if has_label && !has_summary {
            self.append_label_change(target_id, label);
        }

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

    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) {
            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.
    pub fn branch_with_summary(
        &mut self,
        branch_from_id: Option<Uuid>,
        summary: String,
        details: Option<BranchSummaryDetails>,
        from_hook: bool,
    ) -> Uuid {
        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(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.
    pub fn append_label_change(&mut self, target_id: Uuid, label: Option<String>) -> Uuid {
        if !self.entries_by_id.contains_key(&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()
    }

    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::*;

    /// No-op summarizer for tests that don't need summarization.
    struct NoOpSummarizer;
    impl Summarizer for NoOpSummarizer {
        fn summarize(
            &self,
            _entries: &[SessionEntryType],
            _custom_instructions: Option<&str>,
            _replace_instructions: bool,
        ) -> std::pin::Pin<
            Box<
                dyn std::future::Future<Output = Result<BranchSummaryResult, SummarizationError>>
                    + Send
                    + 'static,
            >,
        > {
            Box::pin(async {
                Ok(BranchSummaryResult {
                    summary: None,
                    read_files: vec![],
                    modified_files: vec![],
                    aborted: false,
                    error: None,
                })
            })
        }
    }

    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(),
        })
    }

    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,
        }
    }

    // ===================================================================
    // Original tests
    // ===================================================================

    #[test]
    fn test_navigate_to_user_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);

        let result = nav.navigate_tree(
            user_id,
            NavigationOptions::default(),
            None as Option<&NoOpSummarizer>,
            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(root_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);

        let result = nav.navigate_tree(
            assistant_id,
            NavigationOptions::default(),
            None as Option<&NoOpSummarizer>,
            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);

        let result = nav.navigate_tree(
            entry_id,
            NavigationOptions::default(),
            None as Option<&NoOpSummarizer>,
            None,
        );

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

    #[test]
    fn test_collect_entries_for_branch_summary() {
        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();

        let 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",
            ),
            create_message(
                branch_user_id,
                Some(user_id),
                MessageRole::User,
                "Branch User",
            ),
            create_message(
                branch_assistant_id,
                Some(branch_user_id),
                MessageRole::Assistant,
                "Branch Assistant",
            ),
        ];

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

        let result = nav.collect_entries_for_branch_summary(Some(branch_assistant_id), root_id);

        assert_eq!(result.common_ancestor_id, Some(root_id));
        assert_eq!(result.entries.len(), 3);
    }

    #[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"));

        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());

        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);

        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));

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

    // ===================================================================
    // Additional comprehensive tests
    // ===================================================================

    #[test]
    fn test_new_navigator_has_no_leaf() {
        let nav = SessionNavigator::new();
        assert!(nav.get_leaf_id().is_none());
    }

    #[test]
    fn test_default_navigator_has_no_leaf() {
        let nav = SessionNavigator::default();
        assert!(nav.get_leaf_id().is_none());
    }

    #[test]
    fn test_get_entry_returns_none_for_unknown() {
        let nav = SessionNavigator::new();
        assert!(nav.get_entry(Uuid::new_v4()).is_none());
    }

    #[test]
    fn test_get_branch_returns_empty_when_no_entries() {
        let nav = SessionNavigator::new();
        assert!(nav.get_branch(None).is_empty());
    }

    #[test]
    fn test_get_branch_returns_full_path() {
        let mut nav = SessionNavigator::new();
        let root_id = Uuid::new_v4();
        let mid_id = Uuid::new_v4();
        let leaf_id = Uuid::new_v4();

        nav.add_entry(create_message(root_id, None, MessageRole::User, "Root"));
        nav.add_entry(create_message(
            mid_id,
            Some(root_id),
            MessageRole::Assistant,
            "Mid",
        ));
        nav.add_entry(create_message(
            leaf_id,
            Some(mid_id),
            MessageRole::User,
            "Leaf",
        ));
        nav.branch(leaf_id);

        let branch = nav.get_branch(None);
        assert_eq!(branch.len(), 3);
        assert_eq!(entry_id(branch[0]), root_id);
        assert_eq!(entry_id(branch[1]), mid_id);
        assert_eq!(entry_id(branch[2]), leaf_id);
    }

    #[test]
    fn test_get_children() {
        let mut nav = SessionNavigator::new();
        let parent_id = Uuid::new_v4();
        let child_a = Uuid::new_v4();
        let child_b = Uuid::new_v4();

        nav.add_entry(create_message(parent_id, None, MessageRole::User, "Parent"));
        nav.add_entry(create_message(
            child_a,
            Some(parent_id),
            MessageRole::Assistant,
            "A",
        ));
        nav.add_entry(create_message(
            child_b,
            Some(parent_id),
            MessageRole::Assistant,
            "B",
        ));

        let children = nav.get_children(parent_id);
        assert_eq!(children.len(), 2);
    }

    #[test]
    fn test_get_children_of_leaf() {
        let mut nav = SessionNavigator::new();
        let id = Uuid::new_v4();
        nav.add_entry(create_message(id, None, MessageRole::User, "Solo"));

        let children = nav.get_children(id);
        assert!(children.is_empty());
    }

    #[test]
    fn test_branch_switches_leaf() {
        let mut nav = SessionNavigator::new();
        let id = Uuid::new_v4();
        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));

        nav.branch(id);
        assert_eq!(nav.get_leaf_id(), Some(id));

        nav.reset_leaf();
        assert_eq!(nav.get_leaf_id(), None);
    }

    #[test]
    fn test_reset_leaf() {
        let mut nav = SessionNavigator::new();
        let id = Uuid::new_v4();
        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
        nav.branch(id);
        assert_eq!(nav.get_leaf_id(), Some(id));

        nav.reset_leaf();
        assert!(nav.get_leaf_id().is_none());
    }

    #[test]
    fn test_from_entries_preserves_leaf() {
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let entries = vec![
            create_message(id1, None, MessageRole::User, "A"),
            create_message(id2, Some(id1), MessageRole::Assistant, "B"),
        ];
        let nav = SessionNavigator::from_entries(entries, Some(id2));
        assert_eq!(nav.get_leaf_id(), Some(id2));
        assert!(nav.get_entry(id1).is_some());
        assert!(nav.get_entry(id2).is_some());
    }

    #[test]
    fn test_navigate_to_nonexistent_returns_cancelled() {
        let mut nav = SessionNavigator::new();
        let id = Uuid::new_v4();
        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
        nav.branch(id);

        let result = nav.navigate_tree(
            Uuid::new_v4(),
            NavigationOptions::default(),
            None as Option<&NoOpSummarizer>,
            None,
        );
        assert!(result.cancelled);
    }

    #[test]
    fn test_navigate_to_root_resets_leaf() {
        let mut nav = SessionNavigator::new();
        let root_id = Uuid::new_v4();
        let child_id = Uuid::new_v4();

        nav.add_entry(create_message(root_id, None, MessageRole::User, "Root"));
        nav.add_entry(create_message(
            child_id,
            Some(root_id),
            MessageRole::Assistant,
            "Child",
        ));
        nav.branch(child_id);

        let result = nav.navigate_tree(
            root_id,
            NavigationOptions::default(),
            None as Option<&NoOpSummarizer>,
            None,
        );
        assert!(!result.cancelled);
        assert_eq!(result.editor_text, Some("Root".to_string()));
        assert_eq!(nav.get_leaf_id(), None);
    }

    #[test]
    fn test_collect_entries_no_old_leaf() {
        let target_id = Uuid::new_v4();
        let mut nav = SessionNavigator::new();
        nav.add_entry(create_message(target_id, None, MessageRole::User, "T"));

        let result = nav.collect_entries_for_branch_summary(None, target_id);
        assert!(result.entries.is_empty());
        assert_eq!(result.common_ancestor_id, None);
    }

    #[test]
    fn test_collect_entries_same_branch_common_ancestor() {
        let root_id = Uuid::new_v4();
        let user_id = Uuid::new_v4();
        let assistant_id = Uuid::new_v4();

        let 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, "Asst"),
        ];

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

        let result = nav.collect_entries_for_branch_summary(Some(assistant_id), user_id);
        assert_eq!(result.common_ancestor_id, Some(user_id));
        assert_eq!(result.entries.len(), 1);
    }

    #[test]
    fn test_label_timestamp() {
        let mut nav = SessionNavigator::new();
        let id = Uuid::new_v4();
        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));

        assert!(nav.get_label_timestamp(id).is_none());
        nav.append_label_change(id, Some("marker".to_string()));
        assert!(nav.get_label_timestamp(id).is_some());
    }

    #[test]
    fn test_label_replace() {
        let mut nav = SessionNavigator::new();
        let id = Uuid::new_v4();
        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));

        nav.append_label_change(id, Some("first".to_string()));
        assert_eq!(nav.get_label(id), Some("first"));

        nav.append_label_change(id, Some("second".to_string()));
        assert_eq!(nav.get_label(id), Some("second"));
    }

    #[test]
    fn test_label_nonexistent_entry_returns_nil() {
        let mut nav = SessionNavigator::new();
        let id = nav.append_label_change(Uuid::new_v4(), Some("ghost".to_string()));
        assert_eq!(id, Uuid::nil());
    }

    #[test]
    fn test_branch_with_summary_details() {
        let mut nav = SessionNavigator::new();
        let id = Uuid::new_v4();
        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
        nav.branch(id);

        let details = BranchSummaryDetails {
            read_files: vec!["a.rs".into()],
            modified_files: vec!["b.rs".into()],
        };
        let summary_id = nav.branch_with_summary(Some(id), "Summary".into(), Some(details), true);

        match nav.get_entry(summary_id) {
            Some(SessionEntryType::BranchSummary(e)) => {
                assert_eq!(e.summary, "Summary");
                assert!(e.from_hook.unwrap_or(false));
                assert!(e.details.is_some());
                let d = e.details.as_ref().unwrap();
                assert_eq!(d.read_files, vec!["a.rs"]);
                assert_eq!(d.modified_files, vec!["b.rs"]);
            }
            _ => panic!("Expected branch summary"),
        }
    }

    #[test]
    fn test_navigate_with_extension_cancel() {
        let mut nav = SessionNavigator::new();
        let root_id = Uuid::new_v4();
        let child_id = Uuid::new_v4();
        nav.add_entry(create_message(root_id, None, MessageRole::User, "R"));
        nav.add_entry(create_message(
            child_id,
            Some(root_id),
            MessageRole::Assistant,
            "C",
        ));
        nav.branch(child_id);

        let hook = |_: TreePreparation| -> BeforeTreeHookResult {
            BeforeTreeHookResult {
                cancel: true,
                summary: None,
                custom_instructions: None,
                replace_instructions: None,
                label: None,
            }
        };

        let result = nav.navigate_tree(
            root_id,
            NavigationOptions::default(),
            None as Option<&NoOpSummarizer>,
            Some(&hook),
        );
        assert!(result.cancelled);
    }

    #[test]
    fn test_navigate_with_extension_summary() {
        let mut nav = SessionNavigator::new();
        let root_id = Uuid::new_v4();
        let child_id = Uuid::new_v4();
        nav.add_entry(create_message(root_id, None, MessageRole::User, "R"));
        nav.add_entry(create_message(
            child_id,
            Some(root_id),
            MessageRole::Assistant,
            "C",
        ));
        nav.branch(child_id);

        let hook = |_: TreePreparation| -> BeforeTreeHookResult {
            BeforeTreeHookResult {
                cancel: false,
                summary: Some(ExtensionSummary {
                    summary: "Ext summary".into(),
                    details: None,
                }),
                custom_instructions: None,
                replace_instructions: None,
                label: None,
            }
        };

        let result = nav.navigate_tree(
            root_id,
            NavigationOptions {
                summarize: true,
                ..Default::default()
            },
            None as Option<&NoOpSummarizer>,
            Some(&hook),
        );
        assert!(!result.cancelled);
        assert!(result.summary_entry_id.is_some());

        let sid = result.summary_entry_id.unwrap();
        match nav.get_entry(sid) {
            Some(SessionEntryType::BranchSummary(e)) => {
                assert_eq!(e.summary, "Ext summary");
            }
            _ => panic!("Expected branch summary from extension"),
        }
    }

    #[test]
    fn test_message_role_checks() {
        assert!(MessageRole::User.is_user());
        assert!(!MessageRole::User.is_assistant());
        assert!(MessageRole::Assistant.is_assistant());
        assert!(!MessageRole::Assistant.is_user());
        assert!(!MessageRole::System.is_user());
        assert!(!MessageRole::Tool.is_user());
    }

    #[test]
    fn test_utility_functions() {
        let user_entry = create_message(Uuid::new_v4(), None, MessageRole::User, "hi");
        let asst_entry = create_message(Uuid::new_v4(), None, MessageRole::Assistant, "yo");
        let sys_entry = create_message(Uuid::new_v4(), None, MessageRole::System, "sys");

        assert!(is_user_message(&user_entry));
        assert!(!is_user_message(&asst_entry));
        assert!(is_assistant_message(&asst_entry));
        assert!(!is_assistant_message(&user_entry));
        assert!(!is_user_message(&sys_entry));
        assert!(!is_assistant_message(&sys_entry));
    }

    #[test]
    fn test_session_entry_type_accessors() {
        let id = Uuid::new_v4();
        let msg = SessionEntryType::Message(MessageEntry {
            id,
            parent_id: None,
            timestamp: 42,
            role: MessageRole::User,
            content: "test".into(),
        });

        match &msg {
            SessionEntryType::Message(e) => {
                assert_eq!(e.id, id);
                assert_eq!(e.parent_id, None);
                assert_eq!(e.timestamp, 42);
                assert_eq!(e.role, MessageRole::User);
                assert_eq!(e.content, "test");
            }
            _ => panic!("Expected Message"),
        }
    }

    #[test]
    fn test_get_all_entries() {
        let mut nav = SessionNavigator::new();
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        nav.add_entry(create_message(id1, None, MessageRole::User, "A"));
        nav.add_entry(create_message(id2, Some(id1), MessageRole::Assistant, "B"));

        let all = nav.get_entries();
        assert_eq!(all.len(), 2);
    }
}