1use 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#[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), TaskStatus::Running => Color::Rgb(255, 183, 77), TaskStatus::Verifying => Color::Rgb(129, 212, 250), TaskStatus::Completed => Color::Rgb(102, 187, 106), TaskStatus::Failed => Color::Rgb(239, 83, 80), TaskStatus::Escalated => Color::Rgb(186, 104, 200), }
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#[derive(Debug, Clone)]
65pub struct TaskNode {
66 pub id: String,
68 pub goal: String,
70 pub status: TaskStatus,
72 pub depth: usize,
74 pub parent_id: Option<String>,
76 pub has_children: bool,
78 pub energy: Option<f32>,
80}
81
82#[derive(Default)]
84pub struct TaskTree {
85 nodes: HashMap<String, TaskNode>,
87 roots: Vec<String>,
89 collapsed: HashSet<String>,
91 visible_tasks: Vec<String>,
93 pub state: ListState,
95 theme: Theme,
97}
98
99impl TaskTree {
100 pub fn new() -> Self {
102 Self::default()
103 }
104
105 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 pub fn populate_from_plan(&mut self, plan: perspt_core::types::TaskPlan) {
127 self.clear();
128
129 let mut depth_map: HashMap<String, usize> = HashMap::new();
131
132 for task in &plan.tasks {
134 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 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 if !self.visible_tasks.is_empty() {
156 self.state.select(Some(0));
157 }
158 }
159
160 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 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 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 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 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 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 pub fn expand_all(&mut self) {
237 self.collapsed.clear();
238 self.rebuild_visible();
239 }
240
241 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 fn rebuild_visible(&mut self) {
253 self.visible_tasks.clear();
254
255 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 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 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; }
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 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 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 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 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 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 let indent = " ".repeat(task.depth);
360 let collapse_indicator = if task.has_children {
361 if self.collapsed.contains(&task.id) {
362 "▶ " } else {
364 "▼ " }
366 } else {
367 " " };
369
370 let icon = task.status.icon();
371 let color = task.status.color();
372 let goal = truncate(&task.goal, 35);
373
374 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 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
425fn 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}