1use bevy_ecs::prelude::*;
4use bevy_ecs::system::SystemParam;
5use crossterm::event::{KeyCode, KeyEvent, KeyModifiers};
6use ratatui::style::Stylize;
7use ratatui::text::{Line, Span};
8use rig_ecs::stage::{Effect, Outcome, Settled, SupersededBy};
9use rig_harness::ask::Ask;
10use rig_harness::pick::{Chosen, Highlighted, PickPreview, Picker};
11use rig_tools::shorten;
12use unicode_width::UnicodeWidthStr;
13
14use ratatui::layout::Constraint;
15
16use super::editor::{Editor, MainEditor};
17use super::pane::{TranscriptPane, TranscriptSource};
18use super::panel::{PanelCanvas, Placement, TuiPanel};
19use super::view::TuiView;
20use super::view::{ItemRow, ScrollTo, TranscriptItem};
21
22const PAGE: usize = 10;
24
25#[derive(Component, Clone, Debug, Default, PartialEq, Eq)]
43pub struct ChoiceKeys(pub Vec<(KeyCode, ChoiceAction)>);
44
45#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47pub enum ChoiceAction {
48 Tick,
50 TickAll,
52 TickNone,
54 More,
58 Answer,
60 Close,
62}
63
64#[derive(Component, Clone, Debug, Default, PartialEq, Eq)]
87pub struct PickerInput {
88 pub(crate) filter: String,
90 pub(crate) selected: usize,
92 pub(crate) ticked: Vec<usize>,
95 pub(crate) scroll: usize,
97 pub(crate) draft: Option<String>,
100 pub(crate) more: bool,
103 pub(crate) touched: bool,
106 pub(crate) picked: Option<(String, String)>,
109}
110
111impl PickerInput {
112 pub fn filter(&self) -> &str {
114 &self.filter
115 }
116
117 pub fn ticked(&self) -> &[usize] {
120 &self.ticked
121 }
122
123 pub fn in_input(&self) -> bool {
126 self.draft.is_some()
127 }
128}
129
130pub(crate) struct Listing<'a> {
132 pub(crate) title: &'a str,
133 pub(crate) body: &'a str,
134 pub(crate) items: Vec<Item<'a>>,
135 pub(crate) free_text: bool,
136 pub(crate) ticks: bool,
138 pub(crate) multiple: bool,
140 pub(crate) selected: usize,
141 pub(crate) edit: Option<&'a str>,
143 pub(crate) keys: &'a [(KeyCode, ChoiceAction)],
145}
146
147pub(crate) struct Item<'a> {
149 pub(crate) id: &'a str,
151 pub(crate) label: &'a str,
152 pub(crate) detail: Option<&'a str>,
154 pub(crate) checked: bool,
156 pub(crate) more: Option<&'a str>,
158}
159
160impl Item<'_> {
161 pub(crate) fn line(&self, width: usize) -> Line<'static> {
165 let Some(detail) = self.detail else {
166 return Line::from(shorten(self.label, width));
167 };
168 let reserve = match detail.width() {
171 short if short <= width / 2 => short,
172 _ => width / 3,
173 };
174 let label = shorten(self.label, width.saturating_sub(reserve + 2));
175 let detail = shorten(detail, width.saturating_sub(label.width() + 2));
176 let gap = width.saturating_sub(label.width() + detail.width());
177 Line::from(vec![
178 Span::from(label),
179 Span::from(" ".repeat(gap)),
180 Span::from(detail).dim(),
181 ])
182 }
183}
184
185impl Listing<'_> {
186 pub(crate) fn more(&self, choosing: &PickerInput) -> Option<&str> {
189 let rows = self.rows(&choosing.filter);
190 match rows.get(choosing.selected) {
191 Some(Row::Item(index)) if choosing.more => self.items.get(*index)?.more,
192 _ => None,
193 }
194 }
195
196 pub(crate) fn rows(&self, filter: &str) -> Vec<Row> {
199 let lower = filter.to_lowercase();
200 let mut rows: Vec<Row> = (self.items.iter().enumerate())
201 .filter(|(_, item)| {
202 let label = item.label.to_lowercase();
203 lower.split_whitespace().all(|word| label.contains(word))
204 })
205 .map(|(index, _)| Row::Item(index))
206 .collect();
207 if self.free_text {
208 rows.push(Row::Typed);
209 }
210 rows
211 }
212}
213
214#[derive(Clone, Copy, Debug, PartialEq, Eq)]
216pub(crate) enum Row {
217 Item(usize),
219 Typed,
221}
222
223#[derive(SystemParam)]
225pub(crate) struct Choices<'w, 's> {
226 pickers: Query<'w, 's, &'static Picker>,
227 asks: Query<'w, 's, &'static Ask, (Without<Outcome>, Without<Settled>, Without<SupersededBy>)>,
228 keys: Query<'w, 's, &'static ChoiceKeys>,
229}
230
231impl Choices<'_, '_> {
232 pub(crate) fn shown(&self, entity: Entity) -> Option<Listing<'_>> {
234 let keys = self.keys.get(entity).map_or(&[][..], |keys| &keys.0);
235 if let Ok(picker) = self.pickers.get(entity) {
236 return Some(Listing {
237 title: &picker.title,
238 body: &picker.body,
239 items: picker
240 .items
241 .iter()
242 .map(|item| Item {
243 id: &item.command,
244 label: &item.label,
245 detail: item.detail.as_deref(),
246 checked: item.checked,
247 more: None,
248 })
249 .collect(),
250 free_text: picker.free_text,
251 ticks: picker.multiple,
252 multiple: picker.multiple,
253 selected: picker.selected,
254 edit: None,
255 keys,
256 });
257 }
258 Some(listing(self.asks.get(entity).ok()?, keys))
259 }
260}
261
262fn listing<'a>(ask: &'a Ask, keys: &'a [(KeyCode, ChoiceAction)]) -> Listing<'a> {
264 Listing {
265 title: &ask.title,
266 body: &ask.body,
267 items: ask
268 .choices
269 .iter()
270 .map(|choice| Item {
271 id: &choice.id,
272 label: &choice.label,
273 detail: None,
274 checked: choice.checked,
275 more: choice.more.as_deref(),
276 })
277 .collect(),
278 free_text: ask.free_text,
279 ticks: ask.multiple || (ask.free_text && !ask.choices.is_empty()),
281 multiple: ask.multiple,
282 selected: 0,
283 edit: ask
286 .initial
287 .as_deref()
288 .or((ask.multiline || ask.inline).then_some("")),
289 keys,
290 }
291}
292
293fn carried(choosing: PickerInput, old: &Ask, new: &Ask) -> PickerInput {
297 let id_at = |index: &usize| old.choices.get(*index).map(|choice| choice.id.as_str());
298 let ticked: Vec<&str> = choosing.ticked.iter().filter_map(id_at).collect();
299 let was = |id: &str| old.choices.iter().any(|choice| choice.id == id);
300 let choices = new.choices.iter().enumerate();
301 let ticked = choices.filter(|(_, choice)| match was(&choice.id) {
302 true => ticked.contains(&choice.id.as_str()),
303 false => new.multiple && choice.checked,
304 });
305 let (old_rows, new_rows) = (
306 listing(old, &[]).rows(&choosing.filter),
307 listing(new, &[]).rows(&choosing.filter),
308 );
309 let row = match old_rows.get(choosing.selected) {
310 Some(Row::Item(index)) => id_at(index)
311 .and_then(|id| new.choices.iter().position(|choice| choice.id == id))
312 .map(Row::Item),
313 other => other.copied(),
314 };
315 let selected = row.and_then(|row| new_rows.iter().position(|known| *known == row));
316 PickerInput {
317 ticked: ticked.map(|(index, _)| index).collect(),
318 selected: selected.unwrap_or_default(),
319 ..choosing
320 }
321}
322
323pub(crate) type Showing<'w, 's> = Query<'w, 's, (Entity, &'static mut PickerInput)>;
326
327pub(crate) fn queue_choices(
333 mut view: ResMut<TuiView>,
334 mut input: Query<&mut Editor, With<MainEditor>>,
335 (pickers, filled): (Query<Entity, Added<Picker>>, Query<(), Changed<Picker>>),
336 asks: Query<(Entity, &Effect, &Ask), Added<Ask>>,
337 (superseded, every_ask): (Query<(&Ask, &SupersededBy)>, Query<&Ask>),
338 mut showing: Showing,
339 choices: Choices,
340 mut commands: Commands,
341) {
342 let state = view.bypass_change_detection();
343 let successor = |entity: &Entity| superseded.get(*entity).ok().map(|(_, by)| by.0);
345 for slot in state.queue.iter_mut().chain(state.inline.iter_mut()) {
346 *slot = successor(slot).unwrap_or(*slot);
347 }
348 let mut seen = Vec::new();
349 state.queue.retain(|entity| {
350 let first = !seen.contains(entity);
351 seen.push(*entity);
352 first
353 });
354 let mut before = None;
355 for (entity, shown) in &showing {
356 before = Some(entity);
357 let old = superseded.get(entity).ok();
358 let new = old.and_then(|(_, by)| Some((by.0, every_ask.get(by.0).ok()?)));
359 if let (Some((old, _)), Some((next, new))) = (old, new)
360 && choices.shown(next).is_some()
361 {
362 commands
363 .entity(entity)
364 .try_remove::<(PickerInput, Highlighted)>();
365 commands
366 .entity(next)
367 .try_insert(carried(shown.clone(), old, new));
368 before = Some(next);
369 }
370 }
371 let inline_before = state.inline.clone();
372 let mut asked: Vec<(Entity, &Effect, &Ask)> = asks.iter().collect();
373 asked.sort_by_key(|(_, effect, _)| effect.id());
374 let open = |entity: &Entity| choices.shown(*entity).is_some();
375 state.queue.retain(open);
376 state.inline.retain(open);
377 let (inline, asked): (Vec<_>, Vec<_>) = asked.into_iter().partition(|(.., ask)| ask.inline);
378 for (entity, ..) in inline {
379 if !state.inline.contains(&entity) && open(&entity) {
380 state.inline.push(entity);
381 }
382 }
383 for entity in pickers.iter().chain(asked.into_iter().map(|(ask, ..)| ask)) {
384 if !state.queue.contains(&entity) && open(&entity) {
385 state.queue.push(entity);
386 }
387 }
388 if state.inline != inline_before {
389 view.set_changed();
390 }
391 let front = view.queue.first().copied();
392 if front == before {
393 let refilled = front.filter(|front| filled.contains(*front));
394 let shown = refilled.and_then(|front| Some((showing.get_mut(front).ok()?, front)));
395 if let Some(((_, mut state), front)) = shown
396 && let Some(listing) = choices.shown(front)
397 {
398 refill(&mut state, &listing);
399 }
400 return;
401 }
402 let mut main = input.single_mut().ok();
403 let mut none = Editor::default();
404 let input = main.as_deref_mut().unwrap_or(&mut none);
405 show_first(&view, &choices, (&showing, input), &mut commands);
406}
407
408fn refill(state: &mut PickerInput, shown: &Listing<'_>) {
412 let rows = shown.rows(&state.filter);
413 let item = |row: &Row| match row {
414 Row::Item(index) => shown.items.get(*index),
415 Row::Typed => None,
416 };
417 let picked = state.picked.as_ref();
418 let by = |same: &dyn Fn(&Item<'_>, &(String, String)) -> bool| {
419 picked.and_then(|picked| {
420 rows.iter()
421 .position(|row| item(row).is_some_and(|item| same(item, picked)))
422 })
423 };
424 let found = by(&|item, (id, _)| !id.is_empty() && item.id == id)
425 .or_else(|| by(&|item, (_, label)| item.label == label));
426 let selected = found.unwrap_or(state.selected);
427 state.selected = selected.min(rows.len().saturating_sub(1));
428 state.ticked.retain(|index| *index < shown.items.len());
429}
430
431pub(crate) fn show_first(
435 view: &TuiView,
436 choices: &Choices,
437 (showing, input): (&Showing, &mut Editor),
438 commands: &mut Commands,
439) {
440 for (entity, shown) in showing {
441 if let Some(draft) = &shown.draft {
442 input.set(draft.clone());
443 }
444 commands
445 .entity(entity)
446 .try_remove::<(PickerInput, Highlighted)>();
447 }
448 let first = view.queue.iter().find_map(|&entity| {
449 let shown = choices.shown(entity)?;
450 let draft = shown.edit.map(|initial| {
451 let draft = input.take();
452 input.set(initial);
453 draft
454 });
455 let items = shown.items.iter().enumerate();
456 let ticked: Vec<usize> = items
457 .filter(|(_, item)| shown.multiple && item.checked)
458 .map(|(index, _)| index)
459 .collect();
460 let state = PickerInput {
461 filter: String::new(),
462 selected: shown.selected.min(shown.items.len().saturating_sub(1)),
463 touched: !ticked.is_empty(),
464 ticked,
465 scroll: 0,
466 draft,
467 more: false,
468 picked: None,
469 };
470 Some((entity, state))
471 });
472 if let Some((entity, state)) = first {
473 commands.entity(entity).try_insert(state);
474 }
475}
476
477pub(crate) fn mark_highlighted(
482 mut showing: Query<(Entity, &mut PickerInput, Option<&Highlighted>)>,
483 filled: Query<(), Changed<Picker>>,
484 choices: Choices,
485 mut commands: Commands,
486) {
487 for (entity, mut state, highlighted) in &mut showing {
488 let Some(shown) = choices.shown(entity).filter(|_| state.draft.is_none()) else {
489 continue;
490 };
491 let rows = shown.rows(&state.filter);
492 let index = match rows.get(state.selected) {
493 Some(Row::Item(index)) => Some(*index),
494 _ => None,
495 };
496 let item = index.and_then(|index| shown.items.get(index));
497 let picked = item.map(|item| (item.id.to_owned(), item.label.to_owned()));
498 if state.picked != picked && !filled.contains(entity) {
499 state.bypass_change_detection().picked = picked;
500 }
501 match index {
502 Some(index) if highlighted != Some(&Highlighted(index)) => {
503 commands.entity(entity).try_insert(Highlighted(index));
504 }
505 None if highlighted.is_some() => {
506 commands.entity(entity).try_remove::<Highlighted>();
507 }
508 _ => {}
509 }
510 }
511}
512
513pub(crate) fn choice_key(
516 key: KeyEvent,
517 choosing: &mut PickerInput,
518 shown: &Listing<'_>,
519) -> Option<Option<Chosen>> {
520 let control = key.modifiers.contains(KeyModifiers::CONTROL);
521 let alt = key.modifiers.contains(KeyModifiers::ALT);
522 let shift = key.modifiers.contains(KeyModifiers::SHIFT);
523 let rows = shown.rows(&choosing.filter);
524 let last = rows.len().saturating_sub(1);
525 let typing =
526 !choosing.filter.is_empty() || matches!(rows.get(choosing.selected), Some(Row::Typed));
527 let bound = shown.keys.iter().find(|(code, _)| {
528 *code == key.code && !control && !alt && !(typing && matches!(code, KeyCode::Char(_)))
529 });
530 if let Some((_, action)) = bound {
531 return choice_action(*action, choosing, shown, &rows);
532 }
533 match key.code {
534 KeyCode::Up if shift => choosing.scroll = choosing.scroll.saturating_sub(1),
535 KeyCode::Down if shift => choosing.scroll = choosing.scroll.saturating_add(1),
536 KeyCode::Esc => return Some(None),
537 KeyCode::Char('c') if control => return Some(None),
538 KeyCode::Char(_) if control || key.modifiers.contains(KeyModifiers::ALT) => {}
539 KeyCode::Char(character) => {
540 choosing.filter.push(character);
541 choosing.selected = 0;
542 }
543 KeyCode::Backspace => {
544 choosing.filter.pop();
545 choosing.selected = 0;
546 }
547 KeyCode::Up => choosing.selected = choosing.selected.saturating_sub(1),
548 KeyCode::Down => choosing.selected = (choosing.selected + 1).min(last),
549 KeyCode::PageUp => choosing.selected = choosing.selected.saturating_sub(PAGE),
550 KeyCode::PageDown => choosing.selected = (choosing.selected + PAGE).min(last),
551 KeyCode::Right => choosing.more = true,
552 KeyCode::Left => choosing.more = false,
553 KeyCode::Tab if shown.ticks => {
554 if tick(choosing, shown, &rows) {
555 choosing.selected = (choosing.selected + 1).min(last);
556 }
557 }
558 KeyCode::Enter => return answer(choosing, shown, &rows),
559 _ => {}
560 }
561 None
562}
563
564fn choice_action(
567 action: ChoiceAction,
568 choosing: &mut PickerInput,
569 shown: &Listing<'_>,
570 rows: &[Row],
571) -> Option<Option<Chosen>> {
572 match action {
573 ChoiceAction::Tick if shown.ticks => {
574 tick(choosing, shown, rows);
575 }
576 ChoiceAction::TickAll if shown.multiple => {
577 let items = rows.iter().filter_map(|row| match row {
578 Row::Item(index) => Some(*index),
579 Row::Typed => None,
580 });
581 choosing.ticked = items.collect();
582 choosing.ticked.sort_unstable();
583 choosing.touched = true;
584 }
585 ChoiceAction::TickNone if shown.ticks => {
586 choosing.ticked.clear();
587 choosing.touched = true;
588 }
589 ChoiceAction::More => choosing.more = !choosing.more,
590 ChoiceAction::Answer => return answer(choosing, shown, rows),
591 ChoiceAction::Close => return Some(None),
592 ChoiceAction::Tick | ChoiceAction::TickAll | ChoiceAction::TickNone => {}
593 }
594 None
595}
596
597fn tick(choosing: &mut PickerInput, shown: &Listing<'_>, rows: &[Row]) -> bool {
599 let Some(Row::Item(index)) = rows.get(choosing.selected) else {
600 return false;
601 };
602 match choosing.ticked.binary_search(index) {
603 Ok(at) => {
604 choosing.ticked.remove(at);
605 }
606 Err(_) if !shown.multiple => choosing.ticked = vec![*index],
607 Err(at) => choosing.ticked.insert(at, *index),
608 }
609 choosing.touched = true;
610 true
611}
612
613fn answer(choosing: &PickerInput, shown: &Listing<'_>, rows: &[Row]) -> Option<Option<Chosen>> {
617 match rows.get(choosing.selected)? {
618 Row::Item(_) if choosing.touched && shown.multiple => {
619 Some(Some(Chosen::items(choosing.ticked.clone())))
620 }
621 Row::Item(index) => Some(Some(Chosen::item(*index))),
622 Row::Typed if choosing.filter.trim().is_empty() => None,
623 Row::Typed => {
624 let chosen = Chosen::items(choosing.ticked.clone());
625 Some(Some(chosen.with_text(choosing.filter.trim())))
626 }
627 }
628}
629
630#[derive(Component)]
632pub(crate) struct PreviewPane;
633
634pub(crate) fn follow_previews(
639 showing: Query<(Entity, &PickerInput, Ref<PickPreview>)>,
640 panes: Query<(Entity, &ChildOf), With<PreviewPane>>,
641 mut commands: Commands,
642) {
643 let previewing = |picker: Entity| {
644 let shown = showing.get(picker).ok();
645 let shown = shown.filter(|(_, state, _)| !state.in_input());
646 shown.filter(|(.., preview)| matches!(**preview, PickPreview::Conversation { .. }))
647 };
648 for (pane, of) in &panes {
649 if previewing(of.parent()).is_none() {
650 commands.entity(pane).despawn();
651 }
652 }
653 for (picker, _, preview) in &showing {
654 let PickPreview::Conversation { conversation, at } = &*preview else {
655 continue;
656 };
657 let pane = panes.iter().find(|(_, of)| of.parent() == picker);
658 if previewing(picker).is_none() || pane.is_some() && !preview.is_changed() {
659 continue;
660 }
661 let source = TranscriptSource::Value(conversation.clone());
662 let shown = TranscriptPane::new(source).titled("preview");
663 let pane = match pane {
664 Some((pane, _)) => commands.entity(pane).insert(shown).id(),
665 None => {
666 let within = Placement::Within {
667 panel: picker,
668 part: 1,
669 of: 2,
670 };
671 let spawned = (shown, TuiPanel::new(within), PreviewPane, ChildOf(picker));
672 commands.spawn(spawned).id()
673 }
674 };
675 let item = at.map(TranscriptItem::Message);
676 commands.trigger(ScrollTo {
677 pane,
678 item,
679 row: ItemRow::Top(0),
680 });
681 }
682}
683
684#[derive(Component, Clone, Copy, Debug, Default)]
689#[require(TuiPanel = TuiPanel::new(Placement::Bottom(Constraint::Length(1))))]
690pub struct InlineStrip;
691
692pub(crate) fn draw_inline_strips(
694 mut strips: Query<&mut PanelCanvas, With<InlineStrip>>,
695 view: Res<TuiView>,
696 choices: Choices,
697) {
698 let Some(ask) = view.inline.first().and_then(|&ask| choices.shown(ask)) else {
699 return;
700 };
701 for mut canvas in &mut strips {
702 canvas.render(inline_line(&ask));
703 }
704}
705
706fn inline_line(ask: &Listing<'_>) -> Line<'static> {
709 let mut spans = vec![
710 Span::from("? ").yellow().bold(),
711 Span::from(ask.title.to_owned()).bold(),
712 Span::from(" "),
713 ];
714 let typed = ask.free_text.then_some("Answer…");
715 let labels = ask.items.iter().map(|item| item.label).chain(typed);
716 for (index, label) in labels.take(9).enumerate() {
717 spans.push(Span::from(format!(" [Alt+{} {label}]", index + 1)).cyan());
718 }
719 spans.push(Span::from(" [Alt+0 Dismiss]").dim());
720 Line::from(spans)
721}
722
723#[cfg(test)]
724mod tests;