PluralKit/lib/commands/src/lib.rs

#![feature(let_chains)]

use core::panic;
use std::{cmp::Ordering, collections::HashMap};

uniffi::include_scaffolding!("commands");

mod string;
mod token;
use token::*;

use command_system_macros::commands;

// todo!: move all this stuff into a different file
// lib.rs should just have exported symbols and command definitions

#[derive(Debug, Clone)]
struct TreeBranch {
    current_command_key: Option<String>,
    /// branches.keys(), but sorted by specificity
    possible_tokens: Vec<Token>,
    branches: HashMap<Token, TreeBranch>,
}

impl TreeBranch {
    fn register_command(&mut self, command: Command) {
        let mut current_branch = self;
        // iterate over tokens in command
        for token in command.tokens {
            // recursively get or create a sub-branch for each token
            current_branch = current_branch.branches.entry(token).or_insert(TreeBranch {
                current_command_key: None,
                possible_tokens: vec![],
                branches: HashMap::new(),
            })
        }
        // when we're out of tokens, add an Empty branch with the callback and no sub-branches
        current_branch.branches.insert(
            Token::Empty,
            TreeBranch {
                current_command_key: Some(command.cb),
                possible_tokens: vec![],
                branches: HashMap::new(),
            },
        );
    }

    fn sort_tokens(&mut self) {
        for branch in self.branches.values_mut() {
            branch.sort_tokens();
        }

        // put Value tokens at the end
        // i forget exactly how this works
        // todo!: document this before PR mergs
        self.possible_tokens = self
            .branches
            .keys()
            .into_iter()
            .map(|v| v.clone())
            .collect();
        self.possible_tokens.sort_by(|v, _| {
            if matches!(v, Token::Value(_)) {
                Ordering::Greater
            } else {
                Ordering::Less
            }
        });
    }
}

struct Command {
    tokens: Vec<Token>,
    help: String,
    cb: String,
}

// todo: aliases
// todo: categories
commands! {
    command(help, help, "Shows the help command");

    command(member new, member_new, "Creates a new system member");
    command(member @MemberRef, member_show, "Shows information about a member");
    command(member @MemberRef description, member_desc_show, "Shows a member's description");
    command(member @MemberRef description @FullString, member_desc_update, "Changes a member's description");
    command(member @MemberRef privacy, member_privacy_show, "Displays a member's current privacy settings");
    command(member @MemberRef privacy @MemberPrivacyTarget @PrivacyLevel, member_privacy_update, "Changes a member's privacy settings");
}

pub enum CommandResult {
    Ok { command: ParsedCommand },
    Err { error: String },
}

pub struct ParsedCommand {
    pub command_ref: String,
    pub args: Vec<String>,
    pub flags: HashMap<String, Option<String>>,
}

/// Find the next token from an either raw or partially parsed command string
///
/// Returns:
/// - matched token, to move deeper into the tree
/// - matched value (if this command matched an user-provided value such as a member name)
/// - end position of matched token
/// - optionally a short-circuit error
fn next_token(
    possible_tokens: Vec<Token>,
    input: String,
    current_pos: usize,
) -> Result<(Token, Option<String>, usize), Option<String>> {
    // get next parameter, matching quotes
    let param = crate::string::next_param(input.clone(), current_pos);
    println!("matched: {param:?}\n---");

    // try checking if this is a flag
    // todo!: this breaks full text matching if the full text starts with a flag
    // (but that's kinda already broken anyway)
    if let Some((value, new_pos)) = param.clone()
        && value.starts_with('-')
    {
        return Ok((
            Token::Flag,
            Some(value.trim_start_matches('-').to_string()),
            new_pos,
        ));
    }

    // iterate over tokens and run try_match
    for token in possible_tokens {
        if let TokenMatchResult::Match(value) =
            // for FullString just send the whole string
            token.try_match(if matches!(token, Token::FullString) {
                    if input.is_empty() {
                        None
                    } else {
                        Some(input.clone())
                    }
                } else {
                    param.clone().map(|v| v.0)
                })
        {
            return Ok((token, value, param.map(|v| v.1).unwrap_or(current_pos)));
        }
    }

    Err(None)
}

fn parse_command(input: String) -> CommandResult {
    let mut local_tree: TreeBranch = COMMAND_TREE.clone();

    // end position of all currently matched tokens
    let mut current_pos = 0;

    let mut args: Vec<String> = Vec::new();
    let mut flags: HashMap<String, Option<String>> = HashMap::new();

    loop {
        match next_token(
            local_tree.possible_tokens.clone(),
            input.clone(),
            current_pos,
        ) {
            Ok((found_token, arg, new_pos)) => {
                current_pos = new_pos;
                if let Token::Flag = found_token {
                    flags.insert(arg.unwrap(), None);
                    // don't try matching flags as tree elements
                    continue;
                }

                if let Some(arg) = arg {
                    args.push(arg);
                }

                if let Some(next_tree) = local_tree.branches.get(&found_token) {
                    local_tree = next_tree.clone();
                } else {
                    panic!("found token could not match tree, at {input}");
                }
            }
            Err(None) => {
                if let Some(command_ref) = local_tree.current_command_key {
                    return CommandResult::Ok {
                        command: ParsedCommand {
                            command_ref,
                            args,
                            flags,
                        },
                    };
                }
                // todo: check if last token is a common incorrect unquote (multi-member names etc)
                // todo: check if this is a system name in pk;s command
                return CommandResult::Err {
                    error: "Command not found.".to_string(),
                };
            }
            Err(Some(short_circuit)) => {
                return CommandResult::Err {
                    error: short_circuit,
                };
            }
        }
    }
}