open Base

type key =
  | Arrow_down
  | Arrow_left
  | Arrow_right
  | Arrow_up
  | Backspace
  | Ctrl of char
  | Delete
  | End
  | Enter
  | Esc
  | Home
  | Key of char
  | Nul
  | Page_down
  | Page_up
  | Tab

type t = key

let ctrl c = Ctrl c
let key c = Key c

let of_char = function
  | '\000' -> Nul
  | '\008' | '\127' -> Backspace
  | '\009' -> Tab
  | '\010' | '\013' -> Enter
  | '\027' -> Esc
  | c when Char.(c < ' ') -> Char.to_int c + 64 |> Char.of_int_exn |> ctrl
  | c -> Key c

let to_string = function
  | Arrow_down -> "<down>"
  | Arrow_left -> "<left>"
  | Arrow_right -> "<right>"
  | Arrow_up -> "<up>"
  | Backspace -> "<backspace>"
  | Ctrl c -> Printf.sprintf "<c-%c>" c
  | Delete -> "<delete>"
  | End -> "<end>"
  | Enter -> "<cr>"
  | Esc -> "<esc>"
  | Home -> "<home>"
  | Key ' ' -> "<space>"
  | Key c -> String.of_char c
  | Nul -> "<nil>"
  | Page_down -> "<pgdown>"
  | Page_up -> "<pgup>"
  | Tab -> "<tab>"

let stream =
  let step s c =
    let open Sequence.Step in
    let escaped = function
      | 'A' -> Some Arrow_up
      | 'B' -> Some Arrow_down
      | 'C' -> Some Arrow_right
      | 'D' -> Some Arrow_left
      | 'F' -> Some End
      | 'H' -> Some Home
      | _ -> None
    and tilda = function
      | '1' -> Some Home
      | '3' -> Some Delete
      | '4' -> Some End
      | '5' -> Some Page_up
      | '6' -> Some Page_down
      | '7' -> Some Home
      | '8' -> Some End
      | _ -> None
    in
    match (s, c) with
    | `start, Some '\027' -> Skip { state = `esc }
    | `esc, None -> Yield { value = Esc; state = `start }
    | `esc, Some '[' | `escaped, Some 'O' -> Skip { state = `escaped }
    | `escaped, Some i when Char.(i = '1' || ('3' <= i && i <= '8')) ->
        Skip { state = `tilda i }
    | `escaped, c -> (
        match Option.(c >>= escaped) with
        | None -> Skip { state = `state }
        | Some c -> Yield { value = c; state = `start })
    | `tilda i, Some '~' -> (
        match tilda i with
        | None -> Skip { state = `start }
        | Some k -> Yield { value = k; state = `start })
    | `esc, Some _ | `tilda _, _ -> Skip { state = `start }
    | _, None -> Skip { state = `start }
    | _, Some c -> Yield { value = of_char c; state = `start }
  in
  Sequence.unfold_with ~init:`start ~f:step Terminal.char_stream