open Base
open Unix
open Util

type terminal_size = int * int
type state = { tio : terminal_io; size : terminal_size }

let escape ?(prefix = "") ?(args = []) op =
  let open Bytes in
  let args = List.map ~f:Int.to_string args |> String.concat ~sep:";" in
  let lenp = String.length prefix in
  let lena = String.length args in
  let lenb = 3 + lena + lenp in
  let dst = Bytes.create lenb in
  set dst 0 '\x1b';
  set dst 1 '[';
  From_string.blit ~src:prefix ~src_pos:0 ~dst ~dst_pos:2 ~len:lenp;
  From_string.blit ~src:args ~src_pos:0 ~dst ~dst_pos:(2 + lenp) ~len:lena;
  set dst (length dst - 1) op;
  sequence_of_bytes dst

let clear_screen = escape 'J' ~args:[ 2 ]
let clear_to_eol = escape 'K'
let move_cursor x y = escape 'H' ~args:[ x; y ]
let move_down n = escape 'B' ~args:[ n ]
let move_right n = escape 'C' ~args:[ n ]
let query_cursor_pos = escape 'n' ~args:[ 6 ]
let reset_cursor = move_cursor 1 1

let show_cursor show =
  let cmd = if show then 'h' else 'l' in
  escape cmd ~prefix:"?" ~args:[ 25 ]

let input_bytes = Bytes.create 1

let get_char () =
  let syscall () = read stdin input_bytes 0 1 in
  match handle_unix_error syscall () with
  | 0 -> None
  | _ -> Some (Bytes.get input_bytes 0)

let char_stream : char option Sequence.t =
  Sequence.unfold ~init:() ~f:(fun s -> Some (get_char (), s))

let write_seq =
  let syscall s =
    let buf = Sequence.to_list s |> Bytes.of_char_list in
    single_write stdout buf 0 (Bytes.length buf) |> ignore
  in
  handle_unix_error syscall

let write_lines (lines : char Sequence.t Sequence.t) =
  let crnl = Sequence.of_list [ '\r'; '\n' ] in
  let clear s = Sequence.append s clear_to_eol in
  let syscall seq =
    let buf =
      Sequence.(map ~f:clear seq |> intersperse ~sep:crnl |> concat |> to_list)
      |> Bytes.of_char_list
    in
    single_write stdout buf 0 (Bytes.length buf) |> ignore
  in
  handle_unix_error syscall lines

let cmds_to_sequence l = Sequence.(of_list l |> concat)

let restore_screen () =
  cmds_to_sequence [ clear_screen; reset_cursor; show_cursor true ] |> write_seq

let redraw screen (x, y) =
  let pre = cmds_to_sequence [ show_cursor false; reset_cursor ]
  and post = cmds_to_sequence [ move_cursor x y; show_cursor true ] in
  write_seq pre;
  write_lines screen;
  write_seq post

let get_state () = { tio = handle_unix_error tcgetattr stdin; size = (-1, -1) }

let size =
  let query () =
    cmds_to_sequence [ move_right 999; move_down 999; query_cursor_pos ]
    |> write_seq
  and get_reply () =
    Sequence.(
      char_stream
      |> take_while ~f:Option.is_some
      |> Fn.flip drop_eagerly 2 (* Drop escape sequence '<esc>[' *)
      |> map ~f:(Option.value ~default:'R')
      |> take_while ~f:(fun c -> Char.(c <> 'R'))
      |> to_list |> String.of_char_list |> Stdlib.Scanf.sscanf
      |> fun scanner -> scanner "%d;%d" (fun a b -> (a, b)))
  in
  query &> get_reply

let enable_raw_mode tio =
  let syscall () =
    tcsetattr stdin TCSAFLUSH
      {
        tio with
        c_brkint = false;
        c_csize = 8;
        c_echo = false;
        c_echonl = false;
        c_icanon = false;
        c_icrnl = false;
        c_ignbrk = false;
        c_igncr = false;
        c_inlcr = false;
        c_inpck = false;
        c_isig = false;
        c_istrip = false;
        c_ixon = false;
        c_opost = false;
        c_parenb = false;
        c_parmrk = false;
        c_vmin = 0;
        c_vtime = 1;
      }
  in
  handle_unix_error syscall

let restore_status tio =
  let syscall () = tcsetattr stdin TCSAFLUSH tio in
  handle_unix_error syscall

let init () =
  let state = get_state () in
  enable_raw_mode state.tio ();
  restore_status state.tio |> Stdlib.at_exit;
  restore_screen |> Stdlib.at_exit;
  { state with size = size () }