open Base
open Key

type register = char option
type count = int option

type operation =
  | Noop
  | Yank
  | Delete
  | Change
  | Join
  | Paste_before
  | Paste_after
  | Erase_before
  | Erase_after

type scope = Line | To_bol | To_eol | Down | Left | Right | Up

type command =
  | Type of char
  | Simple of Key.t
  | Partial of Key.t
  | Shortcut of register * count * operation
  | Chord of register * count * operation * count * scope

type t = command

let shortcut ?r ?n c = Shortcut (r, n, c)
let chord ?r ?n1 c ?n2 m = Chord (r, n1, c, n2, m)

let i_stream =
  let step s k =
    let open Sequence.Step in
    match (s, k) with
    | `start, Key c -> Yield { value = Type c; state = `start }
    | `start, _ -> Yield { value = Simple k; state = `start }
    | _, _ -> Skip { state = `start }
  in
  Sequence.unfold_with ~init:`start ~f:step Key.stream

let simple_movements =
  [
    Key ' ';
    Key '0';
    Key 'h';
    Key 'j';
    Key 'k';
    Key 'l';
    Key '$';
    Arrow_up;
    Arrow_down;
    Arrow_left;
    Arrow_right;
    Backspace;
    End;
    Home;
  ]

let to_scope = function
  | Key 'j' | Arrow_down -> Down
  | Key 'h' | Arrow_left | Backspace -> Left
  | Key 'l' | Key ' ' | Arrow_right -> Right
  | Key 'k' | Arrow_up -> Up
  | Key '0' | Home -> To_bol
  | Key '$' | End -> To_eol
  | _ -> failwith "Invalid motion."

let is_simple_movement k = List.mem ~equal:Poly.equal simple_movements k

let instant_operation =
  [ Key 'C'; Key 'D'; Key 'J'; Key 'P'; Key 'X'; Key 'p'; Key 'x' ]

let chord_operation = [ Key 'c'; Key 'd'; Key 'y' ]

let to_op = function
  | Key 'J' -> Join
  | Key 'c' | Key 'C' -> Change
  | Key 'd' | Key 'D' -> Delete
  | Key 'y' | Key 'Y' -> Yank
  | Key 'p' -> Paste_after
  | Key 'P' -> Paste_before
  | Key 'x' -> Erase_after
  | Key 'X' -> Erase_before
  | _ -> failwith "Invalid operation in chord."

let is_chord_operation k = List.mem ~equal:Poly.equal chord_operation k
let is_instant_operation k = List.mem ~equal:Poly.equal instant_operation k

let n_stream =
  let step s k =
    let open Sequence.Step in
    match (s, k) with
    | `start, Key '"' -> Yield { value = Partial k; state = `chord_reg_pre }
    (* Register *)
    | `chord_reg_pre, Key c -> Yield { value = Partial k; state = `chord_reg c }
    (* Count (first) *)
    | `start, Key n when Char.('1' <= n && n <= '9') ->
        let n = Char.to_int n - 48 in
        Yield { value = Partial k; state = `chord_fst_n (None, n) }
    | `chord_reg r, Key n when Char.('1' <= n && n <= '9') ->
        let n = Char.to_int n - 48 in
        Yield { value = Partial k; state = `chord_fst_n (Some r, n) }
    | `chord_fst_n (r, m), Key n when Char.('0' <= n && n <= '9') ->
        let n = (10 * m) + Char.to_int n - 48 in
        Yield { value = Partial k; state = `chord_fst_n (r, n) }
    (* Instant operations *)
    | `start, k when is_instant_operation k ->
        Yield { value = shortcut (to_op k); state = `start }
    | `chord_reg r, k when is_instant_operation k ->
        Yield { value = shortcut ~r (to_op k); state = `start }
    | `chord_fst_n (r, n), k when is_instant_operation k ->
        Yield { value = shortcut ?r ~n (to_op k); state = `start }
    (* Chord operation (first) *)
    | `start, k when is_chord_operation k ->
        Yield { value = Partial k; state = `chord_cmd (None, None, to_op k) }
    | `chord_reg r, k when is_chord_operation k ->
        Yield { value = Partial k; state = `chord_cmd (Some r, None, to_op k) }
    | `chord_fst_n (r, n), k when is_chord_operation k ->
        Yield { value = Partial k; state = `chord_cmd (r, Some n, to_op k) }
    (* Count (second) *)
    | `chord_cmd (r, n1, c), Key n when Char.('1' <= n && n <= '9') ->
        let n = Char.to_int n - 48 in
        Yield { value = Partial k; state = `chord_snd_n (r, n1, c, n) }
    | `chord_snd_n (r, n1, c, n2), Key n when Char.('0' <= n && n <= '9') ->
        let n2 = (10 * n2) + Char.to_int n - 48 in
        Yield { value = Partial k; state = `chord_snd_n (r, n1, c, n2) }
    (* Chord operation (second) *)
    | `chord_cmd (r, n, c), k when is_chord_operation k && Poly.(c = to_op k) ->
        Yield { value = chord ?r ?n1:n c Line; state = `start }
    | `chord_snd_n (r, n1, c, n2), k
      when is_chord_operation k && Poly.(c = to_op k) ->
        Yield { value = chord ?r ?n1 c ~n2 Line; state = `start }
    (* Movement *)
    | (`start | `chord_reg _), k when is_simple_movement k ->
        Yield { value = chord Noop (to_scope k); state = `start }
    | `chord_fst_n (_, n), k when is_simple_movement k ->
        Yield { value = chord ~n1:n Noop (to_scope k); state = `start }
    | `chord_cmd (r, n, c), k when is_simple_movement k ->
        Yield { value = chord ?r ?n1:n c (to_scope k); state = `start }
    | `chord_snd_n (r, n1, c, n2), k when is_simple_movement k ->
        Yield { value = chord ?r ?n1 c ~n2 (to_scope k); state = `start }
    (* Catch-all rules *)
    | `start, _ -> Yield { value = Simple k; state = `start }
    | _, _ -> Skip { state = `start }
  in
  Sequence.unfold_with ~init:`start ~f:step Key.stream