Skip to main content

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    Queued,
19    Planning,
20    Pending,
21    Coding,
22    Running,
23    Verifying,
24    Retrying,
25    SheafCheck,
26    Committing,
27    Completed,
28    Failed,
29    Escalated,
30}
31
32impl TaskStatus {
33    pub fn icon(&self) -> &'static str {
34        match self {
35            TaskStatus::Queued => "◇",
36            TaskStatus::Planning => "◈",
37            TaskStatus::Pending => "○",
38            TaskStatus::Coding => "◉",
39            TaskStatus::Running => "◐",
40            TaskStatus::Verifying => "◑",
41            TaskStatus::Retrying => "↻",
42            TaskStatus::SheafCheck => "⊘",
43            TaskStatus::Committing => "⊙",
44            TaskStatus::Completed => "●",
45            TaskStatus::Failed => "✗",
46            TaskStatus::Escalated => "⚠",
47        }
48    }
49
50    pub fn color(&self) -> Color {
51        match self {
52            TaskStatus::Queued => Color::Rgb(158, 158, 158), // Lighter gray
53            TaskStatus::Planning => Color::Rgb(179, 157, 219), // Light purple
54            TaskStatus::Pending => Color::Rgb(120, 144, 156), // Gray
55            TaskStatus::Coding => Color::Rgb(255, 213, 79),  // Yellow
56            TaskStatus::Running => Color::Rgb(255, 183, 77), // Amber
57            TaskStatus::Verifying => Color::Rgb(129, 212, 250), // Light blue
58            TaskStatus::Retrying => Color::Rgb(255, 152, 0), // Orange
59            TaskStatus::SheafCheck => Color::Rgb(77, 208, 225), // Cyan
60            TaskStatus::Committing => Color::Rgb(165, 214, 167), // Light green
61            TaskStatus::Completed => Color::Rgb(102, 187, 106), // Green
62            TaskStatus::Failed => Color::Rgb(239, 83, 80),   // Red
63            TaskStatus::Escalated => Color::Rgb(186, 104, 200), // Purple
64        }
65    }
66}
67
68impl From<perspt_core::NodeStatus> for TaskStatus {
69    fn from(status: perspt_core::NodeStatus) -> Self {
70        match status {
71            perspt_core::NodeStatus::Queued => TaskStatus::Queued,
72            perspt_core::NodeStatus::Planning => TaskStatus::Planning,
73            perspt_core::NodeStatus::Pending => TaskStatus::Pending,
74            perspt_core::NodeStatus::Coding => TaskStatus::Coding,
75            perspt_core::NodeStatus::Running => TaskStatus::Running,
76            perspt_core::NodeStatus::Verifying => TaskStatus::Verifying,
77            perspt_core::NodeStatus::Retrying => TaskStatus::Retrying,
78            perspt_core::NodeStatus::SheafCheck => TaskStatus::SheafCheck,
79            perspt_core::NodeStatus::Committing => TaskStatus::Committing,
80            perspt_core::NodeStatus::Completed => TaskStatus::Completed,
81            perspt_core::NodeStatus::Failed => TaskStatus::Failed,
82            perspt_core::NodeStatus::Escalated => TaskStatus::Escalated,
83        }
84    }
85}
86
87/// A task node for the tree view
88#[derive(Debug, Clone)]
89pub struct TaskNode {
90    /// Unique identifier
91    pub id: String,
92    /// Task goal/description
93    pub goal: String,
94    /// Current status
95    pub status: TaskStatus,
96    /// Depth in tree (for indentation)
97    pub depth: usize,
98    /// Parent node ID (None for root)
99    pub parent_id: Option<String>,
100    /// Whether this node has children
101    pub has_children: bool,
102    /// Lyapunov energy (if available)
103    pub energy: Option<f32>,
104    /// Retry count (incremented on Retrying status)
105    pub retry_count: usize,
106}
107
108/// Task tree viewer state with expand/collapse support
109#[derive(Default)]
110pub struct TaskTree {
111    /// All task nodes indexed by ID
112    nodes: HashMap<String, TaskNode>,
113    /// Root node IDs (top-level tasks)
114    roots: Vec<String>,
115    /// Currently collapsed node IDs
116    collapsed: HashSet<String>,
117    /// Flattened visible list for display
118    visible_tasks: Vec<String>,
119    /// Selection state
120    pub state: ListState,
121    /// Theme for styling
122    theme: Theme,
123}
124
125impl TaskTree {
126    /// Create a new task tree
127    pub fn new() -> Self {
128        Self::default()
129    }
130
131    /// Add a task to the tree (legacy API for compatibility)
132    pub fn add_task(&mut self, id: String, goal: String, depth: usize) {
133        let node = TaskNode {
134            id: id.clone(),
135            goal,
136            status: TaskStatus::Pending,
137            depth,
138            parent_id: None,
139            has_children: false,
140            energy: None,
141            retry_count: 0,
142        };
143
144        if depth == 0 {
145            self.roots.push(id.clone());
146        }
147
148        self.nodes.insert(id, node);
149        self.rebuild_visible();
150    }
151
152    /// Populate tree from TaskPlan using dependency information for tree structure
153    pub fn populate_from_plan(&mut self, plan: perspt_core::types::TaskPlan) {
154        self.clear();
155
156        // Build a map of task ID to dependencies for depth calculation
157        let mut depth_map: HashMap<String, usize> = HashMap::new();
158
159        // First pass: insert all tasks with initial depth based on dependencies
160        for task in &plan.tasks {
161            // Calculate depth: max depth of dependencies + 1, or 0 if no deps
162            let depth = if task.dependencies.is_empty() {
163                0
164            } else {
165                task.dependencies
166                    .iter()
167                    .filter_map(|dep_id| depth_map.get(dep_id))
168                    .max()
169                    .map(|d| d + 1)
170                    .unwrap_or(0)
171            };
172            depth_map.insert(task.id.clone(), depth);
173
174            // Use first dependency as parent (for tree visualization)
175            // This creates a logical parent-child relationship
176            let parent_id = task.dependencies.first().cloned();
177
178            self.add_task_with_parent(task.id.clone(), task.goal.clone(), parent_id, depth);
179        }
180
181        // Selection reset
182        if !self.visible_tasks.is_empty() {
183            self.state.select(Some(0));
184        }
185    }
186
187    /// Add a task with parent relationship
188    pub fn add_task_with_parent(
189        &mut self,
190        id: String,
191        goal: String,
192        parent_id: Option<String>,
193        depth: usize,
194    ) {
195        // Mark parent as having children
196        if let Some(ref pid) = parent_id {
197            if let Some(parent) = self.nodes.get_mut(pid) {
198                parent.has_children = true;
199            }
200        }
201
202        let is_root = parent_id.is_none();
203        let node = TaskNode {
204            id: id.clone(),
205            goal,
206            status: TaskStatus::Pending,
207            depth,
208            parent_id,
209            has_children: false,
210            energy: None,
211            retry_count: 0,
212        };
213
214        if is_root {
215            self.roots.push(id.clone());
216        }
217
218        self.nodes.insert(id, node);
219        self.rebuild_visible();
220    }
221
222    /// Clear all tasks
223    pub fn clear(&mut self) {
224        self.nodes.clear();
225        self.roots.clear();
226        self.collapsed.clear();
227        self.visible_tasks.clear();
228        self.state.select(None);
229    }
230
231    /// Update task status
232    pub fn update_status(&mut self, id: &str, status: TaskStatus) {
233        if let Some(task) = self.nodes.get_mut(id) {
234            if status == TaskStatus::Retrying {
235                task.retry_count += 1;
236            }
237            task.status = status;
238        }
239    }
240
241    /// PSP-5 Phase 8: Add a node or update its status if already present.
242    ///
243    /// Used during resume to pre-populate the tree from persisted state
244    /// without requiring a full TaskPlan.
245    pub fn add_or_update_node(&mut self, id: &str, goal: &str, status: TaskStatus) {
246        if let Some(task) = self.nodes.get_mut(id) {
247            task.status = status;
248        } else {
249            let node = TaskNode {
250                id: id.to_string(),
251                goal: goal.to_string(),
252                status,
253                depth: 0,
254                parent_id: None,
255                has_children: false,
256                energy: None,
257                retry_count: 0,
258            };
259            self.roots.push(id.to_string());
260            self.nodes.insert(id.to_string(), node);
261            self.rebuild_visible();
262        }
263    }
264
265    /// Update task energy
266    pub fn update_energy(&mut self, id: &str, energy: f32) {
267        if let Some(task) = self.nodes.get_mut(id) {
268            task.energy = Some(energy);
269        }
270    }
271
272    /// Toggle collapse state for selected node
273    pub fn toggle_collapse(&mut self) {
274        if let Some(selected) = self.state.selected() {
275            if let Some(id) = self.visible_tasks.get(selected).cloned() {
276                if let Some(node) = self.nodes.get(&id) {
277                    if node.has_children {
278                        if self.collapsed.contains(&id) {
279                            self.collapsed.remove(&id);
280                        } else {
281                            self.collapsed.insert(id);
282                        }
283                        self.rebuild_visible();
284                    }
285                }
286            }
287        }
288    }
289
290    /// Expand all nodes
291    pub fn expand_all(&mut self) {
292        self.collapsed.clear();
293        self.rebuild_visible();
294    }
295
296    /// Collapse all nodes
297    pub fn collapse_all(&mut self) {
298        for (id, node) in &self.nodes {
299            if node.has_children {
300                self.collapsed.insert(id.clone());
301            }
302        }
303        self.rebuild_visible();
304    }
305
306    /// Rebuild the visible task list based on collapse state
307    fn rebuild_visible(&mut self) {
308        self.visible_tasks.clear();
309
310        // Sort tasks by depth for proper tree structure
311        let mut sorted: Vec<_> = self.nodes.values().collect();
312        sorted.sort_by(|a, b| a.depth.cmp(&b.depth).then_with(|| a.id.cmp(&b.id)));
313
314        // Build parent-children map
315        let mut children_map: HashMap<Option<String>, Vec<String>> = HashMap::new();
316        for node in sorted {
317            children_map
318                .entry(node.parent_id.clone())
319                .or_default()
320                .push(node.id.clone());
321        }
322
323        // DFS traversal respecting collapse state
324        fn dfs(
325            node_id: &str,
326            nodes: &HashMap<String, TaskNode>,
327            children_map: &HashMap<Option<String>, Vec<String>>,
328            collapsed: &HashSet<String>,
329            result: &mut Vec<String>,
330        ) {
331            result.push(node_id.to_string());
332
333            if collapsed.contains(node_id) {
334                return; // Skip children if collapsed
335            }
336
337            if let Some(children) = children_map.get(&Some(node_id.to_string())) {
338                for child_id in children {
339                    if nodes.contains_key(child_id) {
340                        dfs(child_id, nodes, children_map, collapsed, result);
341                    }
342                }
343            }
344        }
345
346        // Start from roots
347        if let Some(root_children) = children_map.get(&None) {
348            for root_id in root_children {
349                dfs(
350                    root_id,
351                    &self.nodes,
352                    &children_map,
353                    &self.collapsed,
354                    &mut self.visible_tasks,
355                );
356            }
357        }
358    }
359
360    /// Select next task
361    pub fn next(&mut self) {
362        let len = self.visible_tasks.len();
363        if len == 0 {
364            return;
365        }
366        let i = match self.state.selected() {
367            Some(i) => {
368                if i >= len - 1 {
369                    0
370                } else {
371                    i + 1
372                }
373            }
374            None => 0,
375        };
376        self.state.select(Some(i));
377    }
378
379    /// Select previous task
380    pub fn previous(&mut self) {
381        let len = self.visible_tasks.len();
382        if len == 0 {
383            return;
384        }
385        let i = match self.state.selected() {
386            Some(i) => {
387                if i == 0 {
388                    len - 1
389                } else {
390                    i - 1
391                }
392            }
393            None => 0,
394        };
395        self.state.select(Some(i));
396    }
397
398    /// Get the currently selected task
399    pub fn selected_task(&self) -> Option<&TaskNode> {
400        self.state
401            .selected()
402            .and_then(|i| self.visible_tasks.get(i))
403            .and_then(|id| self.nodes.get(id))
404    }
405
406    /// Render the task tree
407    pub fn render(&mut self, frame: &mut Frame, area: Rect) {
408        let items: Vec<ListItem> = self
409            .visible_tasks
410            .iter()
411            .filter_map(|id| self.nodes.get(id))
412            .map(|task| {
413                // Build tree connector characters
414                let indent = "  ".repeat(task.depth);
415                let collapse_indicator = if task.has_children {
416                    if self.collapsed.contains(&task.id) {
417                        "▶ " // Collapsed
418                    } else {
419                        "▼ " // Expanded
420                    }
421                } else {
422                    "  " // Leaf node
423                };
424
425                let icon = task.status.icon();
426                let color = task.status.color();
427                let goal = truncate(&task.goal, 35);
428
429                // Build styled spans
430                let mut spans = vec![
431                    Span::styled(indent, Style::default().fg(Color::DarkGray)),
432                    Span::styled(collapse_indicator, Style::default().fg(Color::Cyan)),
433                    Span::styled(format!("{} ", icon), Style::default().fg(color)),
434                ];
435
436                // Add energy if available
437                if let Some(energy) = task.energy {
438                    let energy_style = self.theme.energy_style(energy);
439                    spans.push(Span::styled(format!("[{:.2}] ", energy), energy_style));
440                }
441
442                // Add retry count if > 0
443                if task.retry_count > 0 {
444                    spans.push(Span::styled(
445                        format!("↻{} ", task.retry_count),
446                        Style::default().fg(Color::Rgb(255, 152, 0)),
447                    ));
448                }
449
450                spans.push(Span::styled(
451                    format!("{}: ", task.id),
452                    Style::default().fg(color).add_modifier(Modifier::BOLD),
453                ));
454                spans.push(Span::styled(goal, Style::default().fg(Color::White)));
455
456                ListItem::new(Line::from(spans))
457            })
458            .collect();
459
460        let title = format!(
461            "🌳 Task DAG ({} nodes{})",
462            self.visible_tasks.len(),
463            if !self.collapsed.is_empty() {
464                format!(", {} collapsed", self.collapsed.len())
465            } else {
466                String::new()
467            }
468        );
469
470        let list = List::new(items)
471            .block(
472                Block::default()
473                    .title(title)
474                    .borders(Borders::ALL)
475                    .border_style(Style::default().fg(Color::Rgb(96, 125, 139))),
476            )
477            .highlight_style(
478                Style::default()
479                    .bg(Color::Rgb(55, 71, 79))
480                    .add_modifier(Modifier::BOLD),
481            )
482            .highlight_symbol("→ ");
483
484        frame.render_stateful_widget(list, area, &mut self.state);
485    }
486}
487
488/// Truncate a string to max length with ellipsis
489fn truncate(s: &str, max: usize) -> String {
490    if s.chars().count() > max {
491        format!(
492            "{}...",
493            s.chars().take(max.saturating_sub(3)).collect::<String>()
494        )
495    } else {
496        s.to_string()
497    }
498}
499
500#[cfg(test)]
501mod tests {
502    use super::*;
503
504    #[test]
505    fn test_add_tasks() {
506        let mut tree = TaskTree::new();
507        tree.add_task("root".to_string(), "Root task".to_string(), 0);
508        tree.add_task("child1".to_string(), "Child 1".to_string(), 1);
509
510        assert_eq!(tree.nodes.len(), 2);
511        assert_eq!(tree.visible_tasks.len(), 2);
512    }
513
514    #[test]
515    fn test_update_status() {
516        let mut tree = TaskTree::new();
517        tree.add_task("task1".to_string(), "Test".to_string(), 0);
518        tree.update_status("task1", TaskStatus::Running);
519
520        assert_eq!(tree.nodes.get("task1").unwrap().status, TaskStatus::Running);
521    }
522
523    #[test]
524    fn test_navigation() {
525        let mut tree = TaskTree::new();
526        tree.add_task("t1".to_string(), "Task 1".to_string(), 0);
527        tree.add_task("t2".to_string(), "Task 2".to_string(), 0);
528        tree.add_task("t3".to_string(), "Task 3".to_string(), 0);
529
530        assert!(tree.state.selected().is_none());
531
532        tree.next();
533        assert_eq!(tree.state.selected(), Some(0));
534
535        tree.next();
536        assert_eq!(tree.state.selected(), Some(1));
537
538        tree.previous();
539        assert_eq!(tree.state.selected(), Some(0));
540    }
541
542    #[test]
543    fn test_lifecycle_mapping_all_variants() {
544        // Verify all NodeStatus variants map to the expected TaskStatus
545        use perspt_core::NodeStatus;
546        let mappings = vec![
547            (NodeStatus::Queued, TaskStatus::Queued),
548            (NodeStatus::Planning, TaskStatus::Planning),
549            (NodeStatus::Pending, TaskStatus::Pending),
550            (NodeStatus::Coding, TaskStatus::Coding),
551            (NodeStatus::Running, TaskStatus::Running),
552            (NodeStatus::Verifying, TaskStatus::Verifying),
553            (NodeStatus::Retrying, TaskStatus::Retrying),
554            (NodeStatus::SheafCheck, TaskStatus::SheafCheck),
555            (NodeStatus::Committing, TaskStatus::Committing),
556            (NodeStatus::Completed, TaskStatus::Completed),
557            (NodeStatus::Failed, TaskStatus::Failed),
558            (NodeStatus::Escalated, TaskStatus::Escalated),
559        ];
560        for (node_status, expected) in mappings {
561            let result: TaskStatus = node_status.into();
562            assert_eq!(
563                result, expected,
564                "NodeStatus::{:?} should map to TaskStatus::{:?}",
565                node_status, expected
566            );
567        }
568    }
569
570    #[test]
571    fn test_retry_count_increments_on_retrying() {
572        let mut tree = TaskTree::new();
573        tree.add_task("t1".to_string(), "Task".to_string(), 0);
574        assert_eq!(tree.nodes.get("t1").unwrap().retry_count, 0);
575
576        tree.update_status("t1", TaskStatus::Retrying);
577        assert_eq!(tree.nodes.get("t1").unwrap().retry_count, 1);
578
579        tree.update_status("t1", TaskStatus::Verifying);
580        assert_eq!(tree.nodes.get("t1").unwrap().retry_count, 1);
581
582        tree.update_status("t1", TaskStatus::Retrying);
583        assert_eq!(tree.nodes.get("t1").unwrap().retry_count, 2);
584    }
585
586    #[test]
587    fn test_status_icons_and_colors_unique() {
588        let statuses = vec![
589            TaskStatus::Queued,
590            TaskStatus::Planning,
591            TaskStatus::Pending,
592            TaskStatus::Coding,
593            TaskStatus::Running,
594            TaskStatus::Verifying,
595            TaskStatus::Retrying,
596            TaskStatus::SheafCheck,
597            TaskStatus::Committing,
598            TaskStatus::Completed,
599            TaskStatus::Failed,
600            TaskStatus::Escalated,
601        ];
602        // Every status should have a non-empty icon
603        for s in &statuses {
604            assert!(!s.icon().is_empty(), "{:?} should have an icon", s);
605        }
606    }
607}