//! The key engine reads key presses the way vim does. It collects a count
//! (`5j`), a register (`"a`), a sequence (`gg`), an operator and its target
//! (`dab`), or a key and the character after it (`ma`), and hands over a
//! complete [`Typed`] command. It also says what can come next, for the
//! which-key popup ([`KeyEngine::hints`]).

use std::time::Instant;

use super::command::{Bound, CharKey, Command, Operator, Register, Target, TargetBound};
use super::key::{Key, notation};
use super::keymap::{Binding, Tables};

/// The biggest count. `99999j` would only go somewhere silly.
const MAX_COUNT: u32 = 9999;

/// A complete command, with what was typed around it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Typed {
    pub command: Command,
    /// The number typed first, like the 5 in `5j`. Around an operator, both
    /// counts multiply: `2d3w` is 6.
    pub count: Option<u32>,
    /// The register named with `"`, like `a` in `"ayy`.
    pub register: Option<Register>,
}

impl Typed {
    /// How many times to do it: the count, or once.
    pub fn times(&self) -> u32 {
        self.count.unwrap_or(1)
    }
}

/// What a key press led to.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Outcome {
    /// A complete command.
    Run(Typed),
    /// The keys so far start something: waiting for the next key.
    Waiting,
    /// Esc dropped what was being typed.
    Cancelled,
    /// Nothing starts like this. The keys typed (in notation) were dropped.
    Unknown(String),
}

/// What the engine waits for next.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
enum Stage {
    /// A command's keys: none yet, or the start of a sequence like `g`.
    #[default]
    Command,
    /// A register's name, after `"`.
    Register,
    /// The character after a key like `m`.
    Char(CharKey),
    /// An operator's target, like `ab` after `d`. `count` is the one typed
    /// before the operator.
    Target { operator: Operator, operator_keys: Vec<Key>, count: Option<u32> },
    /// The character of a target like `f` (`dfx`), with the counts so far.
    TargetChar { operator: Operator, kind: CharKey, count: Option<u32> },
}

/// Reads key presses into commands. It remembers what's half typed in
/// between.
#[derive(Debug, Default)]
pub struct KeyEngine {
    stage: Stage,
    /// The count being typed: before the command, or before an operator's
    /// target.
    count: Option<u32>,
    register: Option<Register>,
    /// The keys typed so far in this stage.
    keys: Vec<Key>,
    /// Everything typed so far, as vim's `showcmd` shows it: `2"a`, `dg`.
    shown: String,
    /// When the first key of the sequence was pressed (for which-key's delay).
    since: Option<Instant>,
}

/// What can follow the keys typed so far, for the which-key popup.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Hints {
    /// The keys typed and what they lead to: `g · go to`.
    pub title: String,
    pub entries: Vec<Hint>,
}

/// One possible next key.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Hint {
    pub keys: String,
    pub help: String,
    /// Whether the key opens a group of more keys (shown as `+name`).
    pub group: bool,
}

impl KeyEngine {
    /// Reads one key press, using the tables in effect.
    pub fn press(&mut self, key: Key, tables: &Tables<'_>) -> Outcome {
        if key.is_escape() && self.is_pending() {
            self.reset();
            return Outcome::Cancelled;
        }
        self.shown.push_str(&key.to_string());
        match std::mem::take(&mut self.stage) {
            Stage::Command => self.command_key(key, tables),
            Stage::Register => match key.as_char().and_then(Register::named) {
                Some(register) => {
                    self.register = Some(register);
                    Outcome::Waiting
                }
                None => self.unknown(),
            },
            Stage::Char(kind) => match key.as_char() {
                Some(c) => self.run(kind.command(c)),
                None => self.unknown(),
            },
            Stage::Target { operator, operator_keys, count } => {
                self.target_key(key, tables, operator, operator_keys, count)
            }
            Stage::TargetChar { operator, kind, count } => match key.as_char().and_then(|c| kind.target(c)) {
                Some(target) => {
                    self.count = count;
                    self.run(Command::Operate(operator, target))
                }
                None => self.unknown(),
            },
        }
    }

    /// Forgets anything half typed.
    pub fn reset(&mut self) {
        *self = KeyEngine::default();
    }

    /// Whether something is half typed: a count, a register, a sequence…
    pub fn is_pending(&self) -> bool {
        self.stage != Stage::Command
            || !self.keys.is_empty()
            || self.count.is_some()
            || self.register.is_some()
    }

    /// The keys typed so far, like vim's `showcmd`: `2"ad`.
    pub fn shown(&self) -> &str {
        &self.shown
    }

    /// When the sequence being typed started (counts don't count), for
    /// which-key's delay.
    pub fn pending_since(&self) -> Option<Instant> {
        self.since
    }

    /// The first key of the sequence being typed (to show the leader menu
    /// right away).
    pub fn first_key(&self) -> Option<Key> {
        match &self.stage {
            Stage::Target { operator_keys, .. } => operator_keys.first().copied(),
            _ => self.keys.first().copied(),
        }
    }

    /// What can be typed next, for the which-key popup. `None` while nothing
    /// that has follow-ups is half typed.
    pub fn hints(&self, tables: &Tables<'_>) -> Option<Hints> {
        let fixed = |what: &str, choices: &[(&str, &str)]| Hints {
            title: format!("{} · {what}", self.shown),
            entries: choices
                .iter()
                .map(|(keys, help)| Hint { keys: keys.to_string(), help: help.to_string(), group: false })
                .collect(),
        };
        match &self.stage {
            Stage::Register => Some(fixed("register", Register::HINTS)),
            Stage::Char(kind) | Stage::TargetChar { kind, .. } => {
                Some(fixed("then a character", kind.hints()))
            }
            Stage::Command if self.keys.is_empty() => None,
            Stage::Command => {
                let title = match tables.group_name(&self.keys) {
                    Some(name) => format!("{} · {name}", self.shown),
                    None => self.shown.clone(),
                };
                let entries = next_keys(&self.keys, &tables.commands_starting_with(&self.keys), tables);
                Some(Hints { title, entries })
            }
            Stage::Target { operator, operator_keys, .. } => {
                let mut entries = Vec::new();
                if self.keys.is_empty() {
                    let help = "this node, with its branch".to_string();
                    entries.push(Hint { keys: notation(operator_keys), help, group: false });
                }
                let targets = tables.targets_starting_with(&self.keys);
                for target in targets {
                    let next = &target.keys[..=self.keys.len()];
                    let keys = target.keys[self.keys.len()].to_string();
                    if entries.iter().any(|hint| hint.keys == keys) {
                        continue;
                    }
                    let hint = if target.keys.len() == next.len() {
                        Hint { keys, help: target.help.clone(), group: false }
                    } else {
                        let name = tables.target_group_name(next).unwrap_or("more");
                        Hint { keys, help: format!("+{name}"), group: true }
                    };
                    entries.push(hint);
                }
                let group =
                    tables.target_group_name(&self.keys).map(|name| format!(" {name}")).unwrap_or_default();
                Some(Hints { title: format!("{} · {}{group}", self.shown, operator.name()), entries })
            }
        }
    }

    /// A key while waiting for a command.
    fn command_key(&mut self, key: Key, tables: &Tables<'_>) -> Outcome {
        if self.keys.is_empty() {
            if let Some(count) = with_digit(self.count, key) {
                self.count = Some(count);
                return Outcome::Waiting;
            }
            self.since.get_or_insert_with(Instant::now);
            if key.as_char() == Some('"') && self.register.is_none() {
                self.stage = Stage::Register;
                return Outcome::Waiting;
            }
        }
        self.keys.push(key);
        let candidates = tables.commands_starting_with(&self.keys);
        let Some(exact) = candidates.iter().find(|binding| binding.keys == self.keys) else {
            return if candidates.is_empty() { self.unknown() } else { Outcome::Waiting };
        };
        match exact.bound {
            Bound::Command(command) => self.run(command),
            Bound::Operator(operator) => {
                let operator_keys = std::mem::take(&mut self.keys);
                self.stage = Stage::Target { operator, operator_keys, count: self.count.take() };
                Outcome::Waiting
            }
            Bound::WithChar(kind) => {
                self.keys.clear();
                self.stage = Stage::Char(kind);
                Outcome::Waiting
            }
        }
    }

    /// A key while an operator waits for its target.
    fn target_key(
        &mut self,
        key: Key,
        tables: &Tables<'_>,
        operator: Operator,
        operator_keys: Vec<Key>,
        count: Option<u32>,
    ) -> Outcome {
        if self.keys.is_empty()
            && let Some(target_count) = with_digit(self.count, key)
        {
            self.count = Some(target_count);
            self.stage = Stage::Target { operator, operator_keys, count };
            return Outcome::Waiting;
        }
        self.keys.push(key);
        // The operator's own keys again: this node (or line), as in `dd`.
        if self.keys == operator_keys {
            self.count = multiply(count, self.count);
            return self.run(Command::Operate(operator, Target::Current));
        }
        let candidates = tables.targets_starting_with(&self.keys);
        if let Some(exact) = candidates.iter().find(|target| target.keys == self.keys) {
            return match exact.bound {
                TargetBound::Target(target) => {
                    self.count = multiply(count, self.count);
                    self.run(Command::Operate(operator, target))
                }
                TargetBound::WithChar(kind) => {
                    self.keys.clear();
                    let count = multiply(count, self.count.take());
                    self.stage = Stage::TargetChar { operator, kind, count };
                    Outcome::Waiting
                }
            };
        }
        if candidates.is_empty() && !operator_keys.starts_with(&self.keys) {
            return self.unknown();
        }
        self.stage = Stage::Target { operator, operator_keys, count };
        Outcome::Waiting
    }

    fn run(&mut self, command: Command) -> Outcome {
        let typed = Typed { command, count: self.count, register: self.register };
        self.reset();
        Outcome::Run(typed)
    }

    fn unknown(&mut self) -> Outcome {
        let shown = std::mem::take(&mut self.shown);
        self.reset();
        Outcome::Unknown(shown)
    }
}

/// The count with one more digit typed, if `key` is a digit that continues
/// or starts a count (`0` only continues one: alone, it's a key of its own).
fn with_digit(count: Option<u32>, key: Key) -> Option<u32> {
    let digit = key.as_char()?.to_digit(10)?;
    if digit == 0 && count.is_none() {
        return None;
    }
    Some((count.unwrap_or(0) * 10 + digit).min(MAX_COUNT))
}

/// Two counts typed around an operator multiply (`2d3w` deletes 6 words).
fn multiply(before: Option<u32>, after: Option<u32>) -> Option<u32> {
    match (before, after) {
        (None, None) => None,
        _ => Some((before.unwrap_or(1) * after.unwrap_or(1)).min(MAX_COUNT)),
    }
}

/// The keys that can follow `typed`, one entry per next key: what it does,
/// or `+group` when more keys follow.
fn next_keys(typed: &[Key], candidates: &[&Binding], tables: &Tables<'_>) -> Vec<Hint> {
    let mut hints: Vec<Hint> = Vec::new();
    for binding in candidates {
        let Some(next) = binding.keys.get(typed.len()) else { continue };
        let keys = next.to_string();
        if hints.iter().any(|hint| hint.keys == keys) {
            continue;
        }
        let hint = if binding.keys.len() == typed.len() + 1 {
            Hint { keys, help: binding.help.clone(), group: false }
        } else {
            let name = tables.group_name(&binding.keys[..=typed.len()]).unwrap_or("more");
            Hint { keys, help: format!("+{name}"), group: true }
        };
        hints.push(hint);
    }
    hints
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::input::bindings::Context;
    use crate::input::command::Motion;
    use crate::input::key::parse_keys;
    use crate::input::keymap::{DEFAULT_LEADER, Keymap};

    const NORMAL: &[Context] = &[Context::Canvas, Context::Normal, Context::Motions];

    /// Types `text` (vim notation) and returns what the last key led to.
    fn type_keys(engine: &mut KeyEngine, keymap: &Keymap, text: &str) -> Outcome {
        let tables = keymap.tables(NORMAL);
        let mut last = Outcome::Waiting;
        for key in parse_keys(text, keymap.leader()).unwrap() {
            last = engine.press(key, &tables);
        }
        last
    }

    fn run(command: Command, count: Option<u32>, register: Option<Register>) -> Outcome {
        Outcome::Run(Typed { command, count, register })
    }

    #[test]
    fn sequences_counts_and_registers() {
        let keymap = Keymap::new(DEFAULT_LEADER);
        let mut engine = KeyEngine::default();
        assert_eq!(type_keys(&mut engine, &keymap, "g"), Outcome::Waiting);
        assert_eq!(type_keys(&mut engine, &keymap, "g"), run(Command::Move(Motion::First), None, None));
        assert_eq!(type_keys(&mut engine, &keymap, "12j"), run(Command::Move(Motion::Down), Some(12), None));
        assert_eq!(
            type_keys(&mut engine, &keymap, "\"+p"),
            run(Command::PasteAfter, None, Some(Register::Clipboard))
        );
        assert_eq!(
            type_keys(&mut engine, &keymap, "3\"ap"),
            run(Command::PasteAfter, Some(3), Some(Register::Named('a')))
        );
        assert_eq!(type_keys(&mut engine, &keymap, "gq"), Outcome::Unknown("gq".into()));
        assert_eq!(type_keys(&mut engine, &keymap, "z<Esc>"), Outcome::Cancelled);
        assert!(!engine.is_pending());
    }

    #[test]
    fn operators_wait_for_their_target() {
        let keymap = Keymap::new(DEFAULT_LEADER);
        let mut engine = KeyEngine::default();
        let delete = |target, count| run(Command::Operate(Operator::Delete, target), count, None);
        assert_eq!(type_keys(&mut engine, &keymap, "dd"), delete(Target::Current, None));
        assert_eq!(type_keys(&mut engine, &keymap, "dab"), delete(Target::Branch, None));
        assert_eq!(
            type_keys(&mut engine, &keymap, "2d3d"),
            delete(Target::Current, Some(6)),
            "counts multiply"
        );
        assert_eq!(
            type_keys(&mut engine, &keymap, ">>"),
            run(Command::Operate(Operator::Indent, Target::Current), None, None)
        );
        assert_eq!(type_keys(&mut engine, &keymap, "dx"), Outcome::Unknown("dx".into()));
        assert_eq!(type_keys(&mut engine, &keymap, "ma"), run(Command::SetMark('a'), None, None));
        assert_eq!(type_keys(&mut engine, &keymap, "'A"), run(Command::JumpToMark('A'), None, None));
    }

    #[test]
    fn which_key_lists_what_can_follow() {
        let keymap = Keymap::new(DEFAULT_LEADER);
        let tables = keymap.tables(NORMAL);
        let mut engine = KeyEngine::default();
        assert_eq!(engine.hints(&tables), None);

        type_keys(&mut engine, &keymap, "<Space>");
        let hints = engine.hints(&tables).unwrap();
        assert_eq!(hints.title, "<Space> · actions");
        assert!(hints.entries.contains(&Hint { keys: "c".into(), help: "color…".into(), group: false }));
        assert!(hints.entries.contains(&Hint { keys: "v".into(), help: "+show a view".into(), group: true }));
        assert_eq!(engine.first_key(), Some(DEFAULT_LEADER));

        engine.reset();
        type_keys(&mut engine, &keymap, "d");
        let hints = engine.hints(&tables).unwrap();
        assert_eq!(hints.title, "d · delete");
        let keys: Vec<&str> = hints.entries.iter().map(|hint| hint.keys.as_str()).collect();
        assert_eq!(keys, ["d", "a", "i"]);
        assert!(hints.entries[1].group);

        engine.reset();
        type_keys(&mut engine, &keymap, "\"");
        assert_eq!(engine.hints(&tables).unwrap().title, "\" · register");
    }
}