(*
copyright (c) 2013-2014, simon cruanes
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.
*)

(** {1 Pretty-Printing of Boxes} *)

type position = { x:int ; y: int }

let origin = {x=0; y=0;}

let _move pos x y = {x=pos.x + x; y=pos.y + y}
let _add pos1 pos2 = _move pos1 pos2.x pos2.y
let _minus pos1 pos2 = _move pos1 (- pos2.x) (- pos2.y)
let _move_x pos x = _move pos x 0
let _move_y pos y = _move pos 0 y

(** {2 Output: where to print to} *)

module Output = struct
  type t = {
    put_char : position -> char -> unit;
    put_string : position -> string -> unit;
    put_sub_string : position -> string -> int -> int -> unit;
    flush : unit -> unit;
  }

  let put_char out pos c = out.put_char pos c
  let put_string out pos s = out.put_string pos s
  let put_sub_string out pos s s_i s_len = out.put_sub_string pos s s_i s_len

  (** An internal buffer, suitable for writing efficiently, then
      convertable into a list of lines *)
  type buffer = {
    mutable buf_lines : buf_line array;
    mutable buf_len : int;
  }
  and buf_line = {
    mutable bl_str : string;
    mutable bl_len : int;
  }

  let _make_line _ = {bl_str=""; bl_len=0}

  let _ensure_lines buf i =
    if i >= Array.length buf.buf_lines
    then (
      let lines' = Array.init (2 * i + 5) _make_line in
      Array.blit buf.buf_lines 0 lines' 0 buf.buf_len;
      buf.buf_lines <- lines';
    )

  let _ensure_line line i =
    if i >= String.length line.bl_str
    then (
      let str' = String.make (2 * i + 5) ' ' in
      String.blit line.bl_str 0 str' 0 line.bl_len;
      line.bl_str <- str';
    )

  let _buf_put_char buf pos c =
    _ensure_lines buf pos.y;
    _ensure_line buf.buf_lines.(pos.y) pos.x;
    buf.buf_len <- max buf.buf_len (pos.y+1);
    let line = buf.buf_lines.(pos.y) in
    line.bl_str.[pos.x] <- c;
    line.bl_len <- max line.bl_len (pos.x+1)

  let _buf_put_sub_string buf pos s s_i s_len =
    _ensure_lines buf pos.y;
    _ensure_line buf.buf_lines.(pos.y) (pos.x + s_len);
    buf.buf_len <- max buf.buf_len (pos.y+1);
    let line = buf.buf_lines.(pos.y) in
    String.blit s s_i line.bl_str pos.x s_len;
    line.bl_len <- max line.bl_len (pos.x+s_len)

  let _buf_put_string buf pos s =
    _buf_put_sub_string buf pos s 0 (String.length s)

  (* create a new buffer *)
  let make_buffer () =
    let buf = {
      buf_lines = Array.init 16 _make_line;
      buf_len = 0;
    } in
    let buf_out = {
      put_char = _buf_put_char buf;
      put_sub_string = _buf_put_sub_string buf;
      put_string = _buf_put_string buf;
      flush = (fun () -> ());
    } in
    buf, buf_out

  let buf_to_lines ?(indent=0) buf =
    let buffer = Buffer.create (5 + buf.buf_len * 32) in
    for i = 0 to buf.buf_len - 1 do
      for k = 1 to indent do Buffer.add_char buffer ' ' done;
      let line = buf.buf_lines.(i) in
      Buffer.add_substring buffer line.bl_str 0 line.bl_len;
      Buffer.add_char buffer '\n';
    done;
    Buffer.contents buffer

  let buf_output ?(indent=0) oc buf =
    for i = 0 to buf.buf_len - 1 do
      for k = 1 to indent do output_char oc ' '; done;
      let line = buf.buf_lines.(i) in
      output oc line.bl_str 0 line.bl_len;
      output_char oc '\n';
    done
end

(* find [c] in [s], starting at offset [i] *)
let rec _find s c i =
  if i >= String.length s then None
  else if s.[i] = c then Some i
  else _find s c (i+1)

let rec _lines s i k = match _find s '\n' i with
  | None ->
      if i<String.length s then k (String.sub s i (String.length s-i))
  | Some j ->
      let s' = String.sub s i (j-i) in
      k s';
      _lines s (j+1) k

module Box = struct
  type grid_shape =
    | GridNone
    | GridBars

  type 'a shape =
    | Text of string list  (* list of lines *)
    | Frame of 'a
    | Pad of position * 'a (* vertical and horizontal padding *)
    | Grid of grid_shape * 'a array array

  type t = {
    shape : t shape;
    size : position lazy_t;
  }

  let size box = Lazy.force box.size

  let shape b = b.shape

  let _array_foldi f acc a =
    let acc = ref acc in
    Array.iteri (fun i x -> acc := f !acc i x) a;
    !acc

  let _dim_matrix m =
    if Array.length m = 0 then {x=0;y=0}
    else {y=Array.length m; x=Array.length m.(0); }

  let _map_matrix f m =
    Array.map (Array.map f) m

  (* height of a line composed of boxes *)
  let _height_line a =
    _array_foldi
      (fun h i box ->
        let s = size box in
        max h s.y
      ) 0 a

  (* how large is the [i]-th column of [m]? *)
  let _width_column m i =
    let acc = ref 0 in
    for j = 0 to Array.length m - 1 do
      acc := max !acc (size m.(j).(i)).x
    done;
    !acc

  (* from a matrix [m] (line,column), return two arrays [lines] and [columns],
    with [col.(i)] being the start offset of column [i] and
    [lines.(j)] being the start offset of line [j].
    Those arrays have one more slot to indicate the end position.
    @param bars if true, leave space for bars between lines/columns *)
  let _size_matrix ~bars m =
    let dim = _dim_matrix m in
    (* +1 is for keeping room for the vertical/horizontal line/column *)
    let additional_space = if bars then 1 else 0 in
    (* columns *)
    let columns = Array.make (dim.x + 1) 0 in
    for i = 0 to dim.x - 1 do
      columns.(i+1) <- columns.(i) + (_width_column m i) + additional_space
    done;
    (* lines *)
    let lines = Array.make (dim.y + 1) 0 in
    for j = 1 to dim.y do
      lines.(j) <- lines.(j-1) + (_height_line m.(j-1)) + additional_space
    done;
    (* no trailing bars, adjust *)
    columns.(dim.x) <- columns.(dim.x) - additional_space;
    lines.(dim.y) <- lines.(dim.y) - additional_space;
    lines, columns

  let _size = function
    | Text l ->
        let width = List.fold_left
          (fun acc line -> max acc (String.length line)) 0 l
        in
        { x=width; y=List.length l; }
    | Frame t ->
        let {x;y} = size t in
        { x=x+2; y=y+2; }
    | Pad (dim, b') -> 
        let {x;y} = size b' in
        { x=x+2*dim.x; y=y+2*dim.y; }
    | Grid (style,m) ->
        let bars = match style with
          | GridBars -> true
          | GridNone -> false
        in
        let dim = _dim_matrix m in
        let lines, columns = _size_matrix ~bars m in
        { y=lines.(dim.y); x=columns.(dim.x)}

  let _make shape =
    { shape; size=(lazy (_size shape)); }
end

let line s =
  assert (_find s '\n' 0 = None);
  Box._make (Box.Text [s])

let text s =
  let acc = ref [] in
  _lines s 0 (fun x -> acc := x :: !acc);
  Box._make (Box.Text (List.rev !acc))

let lines l =
  assert (List.for_all (fun s -> _find s '\n' 0 = None) l);
  Box._make (Box.Text l)

let int_ x = line (string_of_int x)
let float_ x = line (string_of_float x)
let bool_ x = line (string_of_bool x)

let frame b =
  Box._make (Box.Frame b)

let pad' ~col ~lines b =
  assert (col >=0 || lines >= 0);
  if col=0 && lines=0
    then b
    else Box._make (Box.Pad ({x=col;y=lines}, b))

let pad b = pad' ~col:1 ~lines:1 b

let hpad col b = pad' ~col ~lines:0 b
let vpad lines b = pad' ~col:0 ~lines b

let grid ?(pad=fun b->b) ?(bars=true) m =
  let m = Box._map_matrix pad m in
  Box._make (Box.Grid ((if bars then Box.GridBars else Box.GridNone), m))

let init_grid ?bars ~line ~col f =
  let m = Array.init line (fun j-> Array.init col (fun i -> f ~line:j ~col:i)) in
  grid ?bars m

let vlist ?pad ?bars l =
  let a = Array.of_list l in
  grid ?pad ?bars (Array.map (fun line -> [| line |]) a)

let hlist ?pad ?bars l =
  grid ?pad ?bars [| Array.of_list l |]

let hlist_map ?bars f l = hlist ?bars (List.map f l)
let vlist_map ?bars f l = vlist ?bars (List.map f l)
let grid_map ?bars f m = grid ?bars (Array.map (Array.map f) m)

let grid_text ?(pad=fun x->x) ?bars m =
  grid_map ?bars (fun x -> pad (text x)) m

let transpose m =
  let dim = Box._dim_matrix m in
  Array.init dim.x
    (fun i -> Array.init dim.y (fun j -> m.(j).(i)))

let _write_vline ~out pos n =
  for j=0 to n-1 do
    Output.put_char out (_move_y pos j) '|'
  done

let _write_hline ~out pos n =
  for i=0 to n-1 do
    Output.put_char out (_move_x pos i) '-'
  done

(* render given box on the output, starting with upper left corner
    at the given position. [expected_size] is the size of the
    available surrounding space. [offset] is the offset of the box
    w.r.t the surrounding box *)
let rec _render ?(offset=origin) ?expected_size ~out b pos =
  match Box.shape b with
    | Box.Text l ->
        List.iteri
          (fun i line ->
            Output.put_string out (_move_y pos i) line
          ) l
    | Box.Frame b' ->
        let {x;y} = Box.size b' in
        Output.put_char out pos '+';
        Output.put_char out (_move pos (x+1) (y+1)) '+';
        Output.put_char out (_move pos 0 (y+1)) '+';
        Output.put_char out (_move pos (x+1) 0) '+';
        _write_hline out (_move_x pos 1) x;
        _write_hline out (_move pos 1 (y+1)) x;
        _write_vline out (_move_y pos 1) y;
        _write_vline out (_move pos (x+1) 1) y;
        _render ~out b' (_move pos 1 1)
    | Box.Pad (dim, b') ->
        let expected_size = Box.size b in
        _render ~offset:(_add dim offset) ~expected_size ~out b' (_add pos dim)
    | Box.Grid (style,m) ->
        let dim = Box._dim_matrix m in
        let bars = match style with
          | Box.GridNone -> false
          | Box.GridBars -> true
        in
        let lines, columns = Box._size_matrix ~bars m in

        (* write boxes *)
        for j = 0 to dim.y - 1 do
          for i = 0 to dim.x - 1 do
            let expected_size = {
              x=columns.(i+1)-columns.(i);
              y=lines.(j+1)-lines.(j);
            } in
            let pos' = _move pos (columns.(i)) (lines.(j)) in
            _render ~expected_size ~out m.(j).(i) pos'
          done;
        done;

        let len_hlines, len_vlines = match expected_size with
          | None -> columns.(dim.x), lines.(dim.y)
          | Some {x;y} -> x,y
        in

        (* write frame if needed *)
        begin match style with
        | Box.GridNone -> ()
        | Box.GridBars ->
          for j=1 to dim.y - 1 do
            _write_hline ~out (_move pos (-offset.x) (lines.(j)-1)) len_hlines
          done;
          for i=1 to dim.x - 1 do
            _write_vline ~out (_move pos (columns.(i)-1) (-offset.y)) len_vlines
          done;
          for j=1 to dim.y - 1 do
            for i=1 to dim.x - 1 do
              Output.put_char out (_move pos (columns.(i)-1) (lines.(j)-1)) '+'
            done
          done
        end

let render out b =
  _render ~out b origin

let to_string b =
  let buf, out = Output.make_buffer () in
  render out b;
  Output.buf_to_lines buf

let output ?(indent=0) oc b =
  let buf, out = Output.make_buffer () in
  render out b;
  Output.buf_output ~indent oc buf;
  flush oc