owlen/crates/owlen-tui/src/state/command_palette.rs

use crate::commands;

/// Encapsulates the command-line style palette used in command mode.
///
/// The palette keeps track of the raw buffer, matching suggestions, and the
/// currently highlighted suggestion index.  It contains no terminal-specific
/// logic which makes it straightforward to unit test.
#[derive(Debug, Clone)]
pub struct ModelPaletteEntry {
    pub id: String,
    pub name: String,
    pub provider: String,
}

impl ModelPaletteEntry {
    fn display_name(&self) -> &str {
        if self.name.is_empty() {
            &self.id
        } else {
            &self.name
        }
    }
}

#[derive(Debug, Clone, Default)]
pub struct CommandPalette {
    buffer: String,
    suggestions: Vec<String>,
    selected: usize,
    models: Vec<ModelPaletteEntry>,
    providers: Vec<String>,
}

impl CommandPalette {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn buffer(&self) -> &str {
        &self.buffer
    }

    pub fn suggestions(&self) -> &[String] {
        &self.suggestions
    }

    pub fn selected_index(&self) -> usize {
        self.selected
    }

    pub fn clear(&mut self) {
        self.buffer.clear();
        self.suggestions.clear();
        self.selected = 0;
    }

    pub fn set_buffer(&mut self, value: impl Into<String>) {
        self.buffer = value.into();
        self.refresh_suggestions();
    }

    pub fn push_char(&mut self, ch: char) {
        self.buffer.push(ch);
        self.refresh_suggestions();
    }

    pub fn pop_char(&mut self) {
        self.buffer.pop();
        self.refresh_suggestions();
    }

    pub fn update_dynamic_sources(
        &mut self,
        models: Vec<ModelPaletteEntry>,
        providers: Vec<String>,
    ) {
        self.models = models;
        self.providers = providers;
        self.refresh_suggestions();
    }

    pub fn select_previous(&mut self) {
        if !self.suggestions.is_empty() {
            self.selected = self.selected.saturating_sub(1);
        }
    }

    pub fn select_next(&mut self) {
        if !self.suggestions.is_empty() {
            let max_index = self.suggestions.len().saturating_sub(1);
            self.selected = (self.selected + 1).min(max_index);
        }
    }

    pub fn apply_selected(&mut self) -> Option<String> {
        let selected = self
            .suggestions
            .get(self.selected)
            .cloned()
            .or_else(|| self.suggestions.first().cloned());
        if let Some(value) = selected.clone() {
            self.buffer = value;
            self.refresh_suggestions();
        }
        selected
    }

    pub fn refresh_suggestions(&mut self) {
        let trimmed = self.buffer.trim();
        self.suggestions = self.dynamic_suggestions(trimmed);
        if self.selected >= self.suggestions.len() {
            self.selected = 0;
        }
    }

    pub fn ensure_suggestions(&mut self) {
        if self.suggestions.is_empty() {
            self.refresh_suggestions();
        }
    }

    fn dynamic_suggestions(&self, trimmed: &str) -> Vec<String> {
        if let Some(rest) = trimmed.strip_prefix("model ") {
            let suggestions = self.model_suggestions("model", rest.trim());
            if suggestions.is_empty() {
                commands::suggestions(trimmed)
            } else {
                suggestions
            }
        } else if let Some(rest) = trimmed.strip_prefix("m ") {
            let suggestions = self.model_suggestions("m", rest.trim());
            if suggestions.is_empty() {
                commands::suggestions(trimmed)
            } else {
                suggestions
            }
        } else if let Some(rest) = trimmed.strip_prefix("provider ") {
            let suggestions = self.provider_suggestions("provider", rest.trim());
            if suggestions.is_empty() {
                commands::suggestions(trimmed)
            } else {
                suggestions
            }
        } else {
            commands::suggestions(trimmed)
        }
    }

    fn model_suggestions(&self, keyword: &str, query: &str) -> Vec<String> {
        if query.is_empty() {
            return self
                .models
                .iter()
                .take(15)
                .map(|entry| format!("{keyword} {}", entry.id))
                .collect();
        }

        let mut matches: Vec<(usize, usize, &ModelPaletteEntry)> = self
            .models
            .iter()
            .filter_map(|entry| {
                commands::match_score(entry.id.as_str(), query)
                    .or_else(|| commands::match_score(entry.name.as_str(), query))
                    .or_else(|| {
                        let composite = format!("{} {}", entry.provider, entry.display_name());
                        commands::match_score(composite.as_str(), query)
                    })
                    .map(|score| (score.0, score.1, entry))
            })
            .collect();

        matches.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)).then(a.2.id.cmp(&b.2.id)));
        matches
            .into_iter()
            .take(15)
            .map(|(_, _, entry)| format!("{keyword} {}", entry.id))
            .collect()
    }

    fn provider_suggestions(&self, keyword: &str, query: &str) -> Vec<String> {
        if query.is_empty() {
            return self
                .providers
                .iter()
                .take(15)
                .map(|provider| format!("{keyword} {}", provider))
                .collect();
        }

        let mut matches: Vec<(usize, usize, &String)> = self
            .providers
            .iter()
            .filter_map(|provider| {
                commands::match_score(provider.as_str(), query)
                    .map(|score| (score.0, score.1, provider))
            })
            .collect();

        matches.sort();
        matches
            .into_iter()
            .take(15)
            .map(|(_, _, provider)| format!("{keyword} {}", provider))
            .collect()
    }
}