perspt_tui/
task_tree.rs

1//! Task Tree Component
2//!
3//! Displays the SRBN Task DAG as an interactive tree view with expand/collapse support.
4
5use crate::theme::Theme;
6use ratatui::{
7    layout::Rect,
8    style::{Color, Modifier, Style},
9    text::{Line, Span},
10    widgets::{Block, Borders, List, ListItem, ListState},
11    Frame,
12};
13use std::collections::{HashMap, HashSet};
14
15/// Node status for display
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum TaskStatus {
18    Pending,
19    Running,
20    Verifying,
21    Completed,
22    Failed,
23    Escalated,
24}
25
26impl TaskStatus {
27    pub fn icon(&self) -> &'static str {
28        match self {
29            TaskStatus::Pending => "○",
30            TaskStatus::Running => "◐",
31            TaskStatus::Verifying => "◑",
32            TaskStatus::Completed => "●",
33            TaskStatus::Failed => "✗",
34            TaskStatus::Escalated => "⚠",
35        }
36    }
37
38    pub fn color(&self) -> Color {
39        match self {
40            TaskStatus::Pending => Color::Rgb(120, 144, 156), // Gray
41            TaskStatus::Running => Color::Rgb(255, 183, 77),  // Amber
42            TaskStatus::Verifying => Color::Rgb(129, 212, 250), // Light blue
43            TaskStatus::Completed => Color::Rgb(102, 187, 106), // Green
44            TaskStatus::Failed => Color::Rgb(239, 83, 80),    // Red
45            TaskStatus::Escalated => Color::Rgb(186, 104, 200), // Purple
46        }
47    }
48}
49
50impl From<perspt_core::NodeStatus> for TaskStatus {
51    fn from(status: perspt_core::NodeStatus) -> Self {
52        match status {
53            perspt_core::NodeStatus::Pending => TaskStatus::Pending,
54            perspt_core::NodeStatus::Running => TaskStatus::Running,
55            perspt_core::NodeStatus::Verifying => TaskStatus::Verifying,
56            perspt_core::NodeStatus::Completed => TaskStatus::Completed,
57            perspt_core::NodeStatus::Failed => TaskStatus::Failed,
58            perspt_core::NodeStatus::Escalated => TaskStatus::Escalated,
59        }
60    }
61}
62
63/// A task node for the tree view
64#[derive(Debug, Clone)]
65pub struct TaskNode {
66    /// Unique identifier
67    pub id: String,
68    /// Task goal/description
69    pub goal: String,
70    /// Current status
71    pub status: TaskStatus,
72    /// Depth in tree (for indentation)
73    pub depth: usize,
74    /// Parent node ID (None for root)
75    pub parent_id: Option<String>,
76    /// Whether this node has children
77    pub has_children: bool,
78    /// Lyapunov energy (if available)
79    pub energy: Option<f32>,
80}
81
82/// Task tree viewer state with expand/collapse support
83#[derive(Default)]
84pub struct TaskTree {
85    /// All task nodes indexed by ID
86    nodes: HashMap<String, TaskNode>,
87    /// Root node IDs (top-level tasks)
88    roots: Vec<String>,
89    /// Currently collapsed node IDs
90    collapsed: HashSet<String>,
91    /// Flattened visible list for display
92    visible_tasks: Vec<String>,
93    /// Selection state
94    pub state: ListState,
95    /// Theme for styling
96    theme: Theme,
97}
98
99impl TaskTree {
100    /// Create a new task tree
101    pub fn new() -> Self {
102        Self::default()
103    }
104
105    /// Add a task to the tree (legacy API for compatibility)
106    pub fn add_task(&mut self, id: String, goal: String, depth: usize) {
107        let node = TaskNode {
108            id: id.clone(),
109            goal,
110            status: TaskStatus::Pending,
111            depth,
112            parent_id: None,
113            has_children: false,
114            energy: None,
115        };
116
117        if depth == 0 {
118            self.roots.push(id.clone());
119        }
120
121        self.nodes.insert(id, node);
122        self.rebuild_visible();
123    }
124
125    /// Populate tree from TaskPlan using dependency information for tree structure
126    pub fn populate_from_plan(&mut self, plan: perspt_core::types::TaskPlan) {
127        self.clear();
128
129        // Build a map of task ID to dependencies for depth calculation
130        let mut depth_map: HashMap<String, usize> = HashMap::new();
131
132        // First pass: insert all tasks with initial depth based on dependencies
133        for task in &plan.tasks {
134            // Calculate depth: max depth of dependencies + 1, or 0 if no deps
135            let depth = if task.dependencies.is_empty() {
136                0
137            } else {
138                task.dependencies
139                    .iter()
140                    .filter_map(|dep_id| depth_map.get(dep_id))
141                    .max()
142                    .map(|d| d + 1)
143                    .unwrap_or(0)
144            };
145            depth_map.insert(task.id.clone(), depth);
146
147            // Use first dependency as parent (for tree visualization)
148            // This creates a logical parent-child relationship
149            let parent_id = task.dependencies.first().cloned();
150
151            self.add_task_with_parent(task.id.clone(), task.goal.clone(), parent_id, depth);
152        }
153
154        // Selection reset
155        if !self.visible_tasks.is_empty() {
156            self.state.select(Some(0));
157        }
158    }
159
160    /// Add a task with parent relationship
161    pub fn add_task_with_parent(
162        &mut self,
163        id: String,
164        goal: String,
165        parent_id: Option<String>,
166        depth: usize,
167    ) {
168        // Mark parent as having children
169        if let Some(ref pid) = parent_id {
170            if let Some(parent) = self.nodes.get_mut(pid) {
171                parent.has_children = true;
172            }
173        }
174
175        let is_root = parent_id.is_none();
176        let node = TaskNode {
177            id: id.clone(),
178            goal,
179            status: TaskStatus::Pending,
180            depth,
181            parent_id,
182            has_children: false,
183            energy: None,
184        };
185
186        if is_root {
187            self.roots.push(id.clone());
188        }
189
190        self.nodes.insert(id, node);
191        self.rebuild_visible();
192    }
193
194    /// Clear all tasks
195    pub fn clear(&mut self) {
196        self.nodes.clear();
197        self.roots.clear();
198        self.collapsed.clear();
199        self.visible_tasks.clear();
200        self.state.select(None);
201    }
202
203    /// Update task status
204    pub fn update_status(&mut self, id: &str, status: TaskStatus) {
205        if let Some(task) = self.nodes.get_mut(id) {
206            task.status = status;
207        }
208    }
209
210    /// Update task energy
211    pub fn update_energy(&mut self, id: &str, energy: f32) {
212        if let Some(task) = self.nodes.get_mut(id) {
213            task.energy = Some(energy);
214        }
215    }
216
217    /// Toggle collapse state for selected node
218    pub fn toggle_collapse(&mut self) {
219        if let Some(selected) = self.state.selected() {
220            if let Some(id) = self.visible_tasks.get(selected).cloned() {
221                if let Some(node) = self.nodes.get(&id) {
222                    if node.has_children {
223                        if self.collapsed.contains(&id) {
224                            self.collapsed.remove(&id);
225                        } else {
226                            self.collapsed.insert(id);
227                        }
228                        self.rebuild_visible();
229                    }
230                }
231            }
232        }
233    }
234
235    /// Expand all nodes
236    pub fn expand_all(&mut self) {
237        self.collapsed.clear();
238        self.rebuild_visible();
239    }
240
241    /// Collapse all nodes
242    pub fn collapse_all(&mut self) {
243        for (id, node) in &self.nodes {
244            if node.has_children {
245                self.collapsed.insert(id.clone());
246            }
247        }
248        self.rebuild_visible();
249    }
250
251    /// Rebuild the visible task list based on collapse state
252    fn rebuild_visible(&mut self) {
253        self.visible_tasks.clear();
254
255        // Sort tasks by depth for proper tree structure
256        let mut sorted: Vec<_> = self.nodes.values().collect();
257        sorted.sort_by(|a, b| a.depth.cmp(&b.depth).then_with(|| a.id.cmp(&b.id)));
258
259        // Build parent-children map
260        let mut children_map: HashMap<Option<String>, Vec<String>> = HashMap::new();
261        for node in sorted {
262            children_map
263                .entry(node.parent_id.clone())
264                .or_default()
265                .push(node.id.clone());
266        }
267
268        // DFS traversal respecting collapse state
269        fn dfs(
270            node_id: &str,
271            nodes: &HashMap<String, TaskNode>,
272            children_map: &HashMap<Option<String>, Vec<String>>,
273            collapsed: &HashSet<String>,
274            result: &mut Vec<String>,
275        ) {
276            result.push(node_id.to_string());
277
278            if collapsed.contains(node_id) {
279                return; // Skip children if collapsed
280            }
281
282            if let Some(children) = children_map.get(&Some(node_id.to_string())) {
283                for child_id in children {
284                    if nodes.contains_key(child_id) {
285                        dfs(child_id, nodes, children_map, collapsed, result);
286                    }
287                }
288            }
289        }
290
291        // Start from roots
292        if let Some(root_children) = children_map.get(&None) {
293            for root_id in root_children {
294                dfs(
295                    root_id,
296                    &self.nodes,
297                    &children_map,
298                    &self.collapsed,
299                    &mut self.visible_tasks,
300                );
301            }
302        }
303    }
304
305    /// Select next task
306    pub fn next(&mut self) {
307        let len = self.visible_tasks.len();
308        if len == 0 {
309            return;
310        }
311        let i = match self.state.selected() {
312            Some(i) => {
313                if i >= len - 1 {
314                    0
315                } else {
316                    i + 1
317                }
318            }
319            None => 0,
320        };
321        self.state.select(Some(i));
322    }
323
324    /// Select previous task
325    pub fn previous(&mut self) {
326        let len = self.visible_tasks.len();
327        if len == 0 {
328            return;
329        }
330        let i = match self.state.selected() {
331            Some(i) => {
332                if i == 0 {
333                    len - 1
334                } else {
335                    i - 1
336                }
337            }
338            None => 0,
339        };
340        self.state.select(Some(i));
341    }
342
343    /// Get the currently selected task
344    pub fn selected_task(&self) -> Option<&TaskNode> {
345        self.state
346            .selected()
347            .and_then(|i| self.visible_tasks.get(i))
348            .and_then(|id| self.nodes.get(id))
349    }
350
351    /// Render the task tree
352    pub fn render(&mut self, frame: &mut Frame, area: Rect) {
353        let items: Vec<ListItem> = self
354            .visible_tasks
355            .iter()
356            .filter_map(|id| self.nodes.get(id))
357            .map(|task| {
358                // Build tree connector characters
359                let indent = "  ".repeat(task.depth);
360                let collapse_indicator = if task.has_children {
361                    if self.collapsed.contains(&task.id) {
362                        "▶ " // Collapsed
363                    } else {
364                        "▼ " // Expanded
365                    }
366                } else {
367                    "  " // Leaf node
368                };
369
370                let icon = task.status.icon();
371                let color = task.status.color();
372                let goal = truncate(&task.goal, 35);
373
374                // Build styled spans
375                let mut spans = vec![
376                    Span::styled(indent, Style::default().fg(Color::DarkGray)),
377                    Span::styled(collapse_indicator, Style::default().fg(Color::Cyan)),
378                    Span::styled(format!("{} ", icon), Style::default().fg(color)),
379                ];
380
381                // Add energy if available
382                if let Some(energy) = task.energy {
383                    let energy_style = self.theme.energy_style(energy);
384                    spans.push(Span::styled(format!("[{:.2}] ", energy), energy_style));
385                }
386
387                spans.push(Span::styled(
388                    format!("{}: ", task.id),
389                    Style::default().fg(color).add_modifier(Modifier::BOLD),
390                ));
391                spans.push(Span::styled(goal, Style::default().fg(Color::White)));
392
393                ListItem::new(Line::from(spans))
394            })
395            .collect();
396
397        let title = format!(
398            "🌳 Task DAG ({} nodes{})",
399            self.visible_tasks.len(),
400            if !self.collapsed.is_empty() {
401                format!(", {} collapsed", self.collapsed.len())
402            } else {
403                String::new()
404            }
405        );
406
407        let list = List::new(items)
408            .block(
409                Block::default()
410                    .title(title)
411                    .borders(Borders::ALL)
412                    .border_style(Style::default().fg(Color::Rgb(96, 125, 139))),
413            )
414            .highlight_style(
415                Style::default()
416                    .bg(Color::Rgb(55, 71, 79))
417                    .add_modifier(Modifier::BOLD),
418            )
419            .highlight_symbol("→ ");
420
421        frame.render_stateful_widget(list, area, &mut self.state);
422    }
423}
424
425/// Truncate a string to max length with ellipsis
426fn truncate(s: &str, max: usize) -> String {
427    if s.chars().count() > max {
428        format!(
429            "{}...",
430            s.chars().take(max.saturating_sub(3)).collect::<String>()
431        )
432    } else {
433        s.to_string()
434    }
435}
436
437#[cfg(test)]
438mod tests {
439    use super::*;
440
441    #[test]
442    fn test_add_tasks() {
443        let mut tree = TaskTree::new();
444        tree.add_task("root".to_string(), "Root task".to_string(), 0);
445        tree.add_task("child1".to_string(), "Child 1".to_string(), 1);
446
447        assert_eq!(tree.nodes.len(), 2);
448        assert_eq!(tree.visible_tasks.len(), 2);
449    }
450
451    #[test]
452    fn test_update_status() {
453        let mut tree = TaskTree::new();
454        tree.add_task("task1".to_string(), "Test".to_string(), 0);
455        tree.update_status("task1", TaskStatus::Running);
456
457        assert_eq!(tree.nodes.get("task1").unwrap().status, TaskStatus::Running);
458    }
459
460    #[test]
461    fn test_navigation() {
462        let mut tree = TaskTree::new();
463        tree.add_task("t1".to_string(), "Task 1".to_string(), 0);
464        tree.add_task("t2".to_string(), "Task 2".to_string(), 0);
465        tree.add_task("t3".to_string(), "Task 3".to_string(), 0);
466
467        assert!(tree.state.selected().is_none());
468
469        tree.next();
470        assert_eq!(tree.state.selected(), Some(0));
471
472        tree.next();
473        assert_eq!(tree.state.selected(), Some(1));
474
475        tree.previous();
476        assert_eq!(tree.state.selected(), Some(0));
477    }
478}