(* See the end of the file for the original license of Linenoise. *)

let max_line = 2048

type key =
  | Enter
  | Ctrl_a
  | Ctrl_b
  | Ctrl_c
  | Ctrl_d
  | Ctrl_e
  | Ctrl_f
  | Ctrl_r
  | Ctrl_p
  | Ctrl_g
  | Backspace
  | Escape_sequence
  | Tab
  | Unknown of Uchar.t

let key_of_char c =
  match Char.code c with
  | 1 -> Ctrl_a
  | 2 -> Ctrl_b
  | 3 -> Ctrl_c
  | 4 -> Ctrl_d
  | 5 -> Ctrl_e
  | 6 -> Ctrl_f
  | 7 -> Ctrl_g
  | 16 -> Ctrl_p
  | 18 -> Ctrl_r
  | 9 -> Tab
  | 13 -> Enter
  | 27 -> Escape_sequence
  | 127 -> Backspace
  | _ -> Unknown (Uchar.of_char c)

module State = struct
  type completion = string -> string list

  type t = {
    ifd : Unix.file_descr;
    ofd : Unix.file_descr;
    buf : bytes;
    buf_len : int;
    prompt : bytes;
    plen : int;
    old_pos : int;
    pos : int;
    len : int;
    (* Update later *)
    cols : int;
    old_rows : int;
    old_row_pos : int;
    history_index : int;
    history : string list;
    saved_buf : string;
    read_buf : Bytes.t;
    in_completion : bool;
    completion_idx : int;
    complete : completion option;
  }

  let buf t = Bytes.sub t.buf 0 t.len

  let make ?(in_completion = false) ?(completion_idx = 0) ?complete
      ?(old_pos = 0) ?(pos = 0) ?(len = 0) ?(history = []) ?(ifd = Unix.stdin)
      ?(ofd = Unix.stdout) ~prompt buf =
    {
      in_completion;
      ifd;
      ofd;
      buf;
      buf_len = Bytes.length buf;
      prompt;
      plen = Bytes.length prompt;
      old_pos;
      pos;
      len;
      cols = 0;
      old_row_pos = 1;
      history;
      old_rows = 0;
      history_index = -1;
      saved_buf = "";
      complete;
      completion_idx;
      read_buf = Bytes.make 1 '\000' (* For reading a character *);
    }

  let override ?in_completion ?completion_idx ?complete ?ifd ?ofd ?buf ?buf_len
      ?prompt ?plen ?old_pos ?pos ?len ?cols ?old_rows ?old_row_pos
      ?history_index ?history ?saved_buf (t : t) =
    let () =
      match buf with
      | None -> ()
      | Some buf -> Bytes.blit buf 0 t.buf 0 (Bytes.length buf)
    in
    {
      in_completion = Option.value ~default:t.in_completion in_completion;
      ifd = Option.value ~default:t.ifd ifd;
      ofd = Option.value ~default:t.ofd ofd;
      buf = t.buf;
      buf_len = Option.value ~default:t.buf_len buf_len;
      prompt = Option.value ~default:t.prompt prompt;
      plen = Option.value ~default:t.plen plen;
      old_pos = Option.value ~default:t.old_pos old_pos;
      pos = Option.value ~default:t.pos pos;
      len = Option.value ~default:t.len len;
      cols = Option.value ~default:t.cols cols;
      old_rows = Option.value ~default:t.old_rows old_rows;
      old_row_pos = Option.value ~default:t.old_row_pos old_row_pos;
      history_index = Option.value ~default:t.history_index history_index;
      complete = (match complete with Some f -> Some f | None -> t.complete);
      read_buf = t.read_buf;
      completion_idx = Option.value ~default:t.completion_idx completion_idx;
      history = Option.value ~default:t.history history;
      saved_buf = Option.value ~default:t.saved_buf saved_buf;
    }
end

let get_columns () =
  match Terminal.Size.get_columns () with Some n -> n | None -> 80

let with_raw_mode (state : State.t) fn =
  let saved_tio = Unix.tcgetattr state.ifd in
  let tio : Unix.terminal_io =
    {
      saved_tio with
      c_brkint = false;
      c_icrnl = false;
      c_inpck = false;
      c_istrip = false;
      c_ixon = false;
      c_opost = false;
      c_csize = 8;
      c_echo = false;
      c_icanon = false;
      c_isig = false;
      c_vtime = 0;
      c_vmin = 1;
    }
  in
  Unix.tcsetattr state.ifd TCSAFLUSH tio;
  Fun.protect
    ~finally:(fun () -> Unix.tcsetattr state.ifd TCSADRAIN saved_tio)
    fn

let write_bytes fd s =
  let len = Bytes.length s in
  let wrote = Unix.write fd s 0 len in
  assert (Int.equal len wrote)

let write_uchar fd u =
  let b_len = Uchar.utf_8_byte_length u in
  let bs = Bytes.create b_len in
  let wrote = Bytes.set_utf_8_uchar bs 0 u in
  assert (Int.equal b_len wrote);
  write_bytes fd bs

type edit = Editing of State.t | Finished of bytes option | Ctrl_c

let read_char state =
  try
    let read = Unix.read state.State.ifd state.read_buf 0 1 in
    if read = 0 then `None else `Some (Bytes.unsafe_get state.read_buf 0)
  with Unix.Unix_error ((Unix.EWOULDBLOCK | Unix.EAGAIN), _, _) -> `Editing

let edit_start ~stdin:_ ~stdout:_ state fn =
  with_raw_mode state @@ fun () ->
  let cols = get_columns () in
  (* Bytes.set state.buf 0 '\000'; *)
  let state = State.override ~cols ~buf_len:(state.buf_len - 1) state in
  write_bytes state.ofd state.prompt;
  fn state

let utf8_display_width b len =
  let s = Bytes.to_string b in
  Terminal.guess_printed_width (String.sub s 0 len)

let utf8_next_char_len s off =
  Bytes.get_utf_8_uchar s off
  |> Uchar.utf_decode_uchar |> Uchar.utf_8_byte_length

let utf8_prev_char_len s off =
  let rec loop pos =
    if pos < 0 || pos < off - 4 then
      invalid_arg "UTF8 previous character length";
    let decode = Bytes.get_utf_8_uchar s off in
    if Uchar.utf_decode_is_valid decode then
      Uchar.utf_decode_uchar decode |> Uchar.utf_8_byte_length
    else loop (pos - 1)
  in
  loop (off - 1)

type refresh_flag = Rewrite

let refresh_single_line ?(flags = []) ?prompt (state : State.t) =
  let prompt = match prompt with None -> state.prompt | Some p -> p in
  let pwidth = utf8_display_width prompt state.plen in
  let poscol = ref @@ utf8_display_width state.buf state.pos in
  let lencol = ref @@ utf8_display_width state.buf state.len in

  let rec loop (state : State.t) =
    if pwidth + !poscol >= state.cols then begin
      let clen = utf8_next_char_len state.buf 0 in
      let c_width =
        Uchar.utf_8_byte_length
          (Bytes.get_utf_8_uchar state.buf clen |> Uchar.utf_decode_uchar)
      in
      poscol := !poscol - c_width;
      lencol := !lencol - c_width;
      let state =
        State.override ~len:(state.len - clen) ~pos:(state.pos - clen) state
      in
      loop state
    end
    else state
  in
  let state = loop state in
  let ab = Buffer.create 0 in
  (* Clear line *)
  Buffer.add_char ab '\r';

  (* Add prompt *)
  if List.mem Rewrite flags then begin
    Buffer.add_bytes ab prompt;
    Buffer.add_bytes ab (Bytes.sub state.buf 0 state.len)
  end;

  (* Erase to the right *)
  Buffer.add_string ab "\x1b[0K";

  (* Cursor to the original position *)
  if List.mem Rewrite flags then begin
    Buffer.add_string ab (Format.sprintf "\r\x1b[%dC" (!poscol + pwidth))
  end;
  write_bytes state.ofd (Buffer.to_bytes ab);
  state

let refresh_line state = refresh_single_line ~flags:[ Rewrite ] state

let refresh_line_with_completions (state : State.t) lcs =
  if state.completion_idx < List.length lcs then
    let saved_state = state in
    let saved_buf = Bytes.copy state.buf in
    let state =
      State.override
        ~len:(Bytes.length (List.nth lcs state.completion_idx))
        ~pos:(Bytes.length (List.nth lcs state.completion_idx))
        ~buf:(List.nth lcs state.completion_idx)
        state
    in
    let state : State.t = refresh_line state in
    State.override state ~len:saved_state.len ~pos:saved_state.pos
      ~buf:saved_buf
  else state

let edit_insert (state : State.t) c =
  let clen = Uchar.utf_8_byte_length c in
  (* At the end of the line *)
  if Int.equal state.len state.pos then begin
    let _ : int = Bytes.set_utf_8_uchar state.buf state.pos c in
    let state =
      State.override ~pos:(state.pos + clen) ~len:(state.len + clen) state
    in
    if
      utf8_display_width state.prompt state.plen
      + utf8_display_width state.buf state.len
      < state.cols
    then begin
      write_uchar state.ofd c;
      state
    end
    else refresh_line state
  end
  else begin
    Bytes.blit state.buf state.pos state.buf (state.pos + clen)
      (state.len - state.pos);
    let _ : int = Bytes.set_utf_8_uchar state.buf state.pos c in
    let state =
      State.override ~len:(state.len + clen) ~pos:(state.pos + clen) state
    in
    refresh_line state
  end

let edit_backspace (state : State.t) =
  let state =
    if state.pos > 0 && state.len > 0 then begin
      let clen = utf8_prev_char_len state.buf state.pos in
      let dst = state.pos - clen in
      let src = state.pos in
      let len = state.len - state.pos in
      Bytes.blit state.buf src state.buf dst len;
      State.override ~pos:(state.pos - clen) ~len:(state.len - clen) state
    end
    else state
  in
  refresh_line state

let get_bytes state = Bytes.sub state.State.buf 0 state.len

let complete_line (state : State.t) c cn =
  match cn (String.of_bytes (get_bytes state)) with
  | [] -> (State.override ~in_completion:false state, `Char c)
  | xs ->
      let state, c =
        match key_of_char c with
        | Tab ->
            if not state.in_completion then
              ( State.override ~in_completion:true ~completion_idx:0 state,
                `Edit_more )
            else
              ( State.override
                  ~completion_idx:
                    ((state.completion_idx + 1) mod (List.length xs + 1))
                  state,
                `Edit_more )
        | _ ->
            let state =
              if state.completion_idx < List.length xs then begin
                let to_write = List.nth xs state.completion_idx in
                let to_write_len = String.length to_write in
                Bytes.blit_string to_write 0 state.buf 0 to_write_len;
                State.override ~len:to_write_len ~pos:to_write_len state
              end
              else state
            in
            (State.override ~in_completion:false state, `Char c)
      in
      if state.in_completion && state.completion_idx < List.length xs then begin
        (refresh_line_with_completions state (List.map Bytes.of_string xs), c)
      end
      else begin
        (refresh_line state, c)
      end

let move_left (state : State.t) =
  let s =
    if state.pos > 0 then
      State.override
        ~pos:(state.pos - utf8_prev_char_len state.buf state.pos)
        state
    else state
  in
  refresh_line s

let move_right (state : State.t) =
  let s =
    if state.pos > 0 then
      State.override
        ~pos:(state.pos + utf8_next_char_len state.buf state.pos)
        state
    else state
  in
  refresh_line s

let reverse_incr_search ~history (state : State.t) =
  let has_match = ref true in
  let search_buf = Buffer.create 16 in
  let search_pos = ref 0 in
  let search_dir = ref (-1) in
  let h = history "" in
  let history_len = List.length h in
  let saved_buf = Bytes.copy state.buf in
  let exception Completed of State.t in
  let rec loop state : State.t =
    let prompt =
      if !has_match then
        Fmt.str "(reverse-i-search)`%s': " (Buffer.contents search_buf)
      else
        Fmt.str "(failed-reverse-i-search)`%s': " (Buffer.contents search_buf)
    in
    let new_char = ref false in
    let state = State.override ~pos:0 state in
    let state =
      refresh_single_line ~flags:[ Rewrite ] ~prompt:(String.to_bytes prompt)
        state
    in
    let state =
      match read_char state with
      | `Editing -> loop state
      | `None -> loop state
      | `Some c -> (
          match key_of_char c with
          | Backspace ->
              if Buffer.length search_buf > 0 then begin
                (* Pretty wasteful... *)
                let s = Buffer.contents search_buf in
                Buffer.clear search_buf;
                Buffer.add_substring search_buf s 0 (String.length s - 1);
                search_pos := 0
              end;
              state
          | Ctrl_p ->
              search_dir := -1;
              if !search_pos >= history_len then search_pos := history_len - 1;
              state
          | Ctrl_r ->
              search_dir := 1;
              if !search_pos < 0 then search_pos := 0;
              state
          | Ctrl_g ->
              let l = Bytes.length saved_buf in
              Bytes.blit saved_buf 0 state.buf 0 l;
              let state = refresh_line (State.override ~pos:l ~len:l state) in
              raise (Completed state)
          | Enter ->
              let state = State.override ~pos:state.len state in
              raise (Completed state)
          | _ ->
              if Char.compare c ' ' > 0 then begin
                new_char := true;
                Buffer.add_char search_buf c;
                search_pos := 0;
                state
              end
              else
                State.override ~pos:state.len state |> refresh_line |> fun s ->
                raise (Completed s))
    in
    has_match := false;
    let state =
      if Buffer.length search_buf > 0 then begin
        let rec inner_loop () =
          if !search_pos >= 0 && !search_pos < history_len then begin
            let entry = List.nth h !search_pos in
            match
              ( Astring.String.cut ~sep:(Buffer.contents search_buf) entry,
                !new_char
                || not
                   @@ String.equal entry (Bytes.to_string (State.buf state)) )
            with
            | Some (_l, _r), true ->
                has_match := true;
                Bytes.blit_string entry 0 state.buf 0 (String.length entry);
                let state = State.override ~len:(String.length entry) state in
                state
            | _ ->
                search_pos := !search_pos + !search_dir;
                inner_loop ()
          end
          else state
        in
        inner_loop ()
      end
      else state
    in
    loop state
  in
  try loop state with Completed state -> state

let edit_history dir fn (state : State.t) =
  let saved_state = state in
  let current_buf = Bytes.sub_string state.buf 0 state.len in
  let state =
    match (dir, state.history_index) with
    | `Prev, -1 ->
        State.override ~history:(fn current_buf) ~history_index:0
          ~saved_buf:current_buf state
    | `Prev, m ->
        let max_history = List.length state.history in
        if m < max_history - 1 then
          State.override ~history_index:(state.history_index + 1) state
        else state
    | `Next, m when m >= 0 ->
        State.override ~history_index:(state.history_index - 1) state
    | _ -> state
  in
  match (state.history, state.history_index) with
  | [], _ -> saved_state
  | _, -1 ->
      let len = String.length state.saved_buf in
      State.override ~buf:(Bytes.of_string state.saved_buf) ~pos:len ~len state
      |> refresh_line
  | _ ->
      let max_history = List.length state.history in
      let idx = min max_history state.history_index in
      let s = List.nth state.history idx in
      let s_len = String.length s in
      State.override ~buf:(Bytes.of_string s) ~pos:s_len ~len:s_len state
      |> refresh_line

let edit_feed ~history state =
  match read_char state with
  | `Editing -> Editing state
  | `None -> Finished None
  | `Some c -> (
      let uc = Uchar.of_char c in
      let state, c =
        if state.in_completion || key_of_char c = Tab then
          match state.complete with
          | None -> (state, `Char c)
          | Some cn -> complete_line state c cn
        else (state, `Char c)
      in
      match c with
      | `Edit_more -> Editing state
      | `Char c -> (
          match key_of_char c with
          | Enter -> Finished (Some (Bytes.sub state.buf 0 state.len))
          | Ctrl_d ->
              if Int.equal state.len 0 then Finished None else Editing state
          | Ctrl_c -> Ctrl_c
          | Ctrl_b -> Editing (move_left state)
          | Ctrl_f -> Editing (move_right state)
          | Ctrl_r -> Editing (reverse_incr_search ~history state)
          | Backspace -> Editing (edit_backspace state)
          | Tab -> Editing state
          | Escape_sequence -> (
              let c1 =
                read_char state |> function `Some c -> c | _ -> assert false
              in
              let c2 =
                read_char state |> function `Some c -> c | _ -> assert false
              in
              match (c1, c2) with
              | '[', 'A' -> Editing (edit_history `Prev history state)
              | '[', 'B' -> Editing (edit_history `Next history state)
              | '[', 'C' -> Editing (move_right state)
              | '[', 'D' -> Editing (move_left state)
              | _ -> Editing state)
          | Unknown _ | _ ->
              let state = edit_insert state uc in
              Editing state))

type result = String of string option | Ctrl_c

let blocking_edit ?complete ~history ~stdin ~stdout buf ~prompt =
  let state = State.make ?complete ~prompt buf in
  let res =
    edit_start ~stdin ~stdout state @@ fun state ->
    let rec loop = function
      | Editing state -> loop (edit_feed ~history state)
      | Finished s -> String (Option.map Bytes.to_string s)
      | Ctrl_c -> Ctrl_c
    in
    loop (edit_feed ~history state)
  in
  res

type history = string -> string list

let bruit ?complete ?(history = fun _ -> []) prompt =
  let prompt = Bytes.of_string prompt in
  let buf = Bytes.make max_line '\000' in
  if not (Unix.isatty Unix.stdin) then failwith "Stdin is not a tty"
  else
    blocking_edit ?complete ~history ~stdin:Unix.stdin ~stdout:Unix.stdout buf
      ~prompt

(*
 * Copyright (c) 2010-2023, Salvatore Sanfilippo <antirez at gmail dot com>
 * Copyright (c) 2010-2013, Pieter Noordhuis <pcnoordhuis at gmail dot com>
 *
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *  *  Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *
 *  *  Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. *)