open Import

module Escape : sig
  val quoted : string -> string
end = struct
  let quote_length s =
    let n = ref 0 in
    let len = String.length s in
    for i = 0 to len - 1 do
      n
      := !n
         +
         match String.unsafe_get s i with
         | '\"' | '\\' | '\n' | '\t' | '\r' | '\b' -> 2
         | '%' -> if i + 1 < len && s.[i + 1] = '{' then 2 else 1
         | ' ' .. '~' -> 1
         | _ -> 4
    done;
    !n
  ;;

  let escape_to s ~dst:s' ~ofs =
    let n = ref ofs in
    let len = String.length s in
    for i = 0 to len - 1 do
      (match String.unsafe_get s i with
       | ('\"' | '\\') as c ->
         Bytes.unsafe_set s' !n '\\';
         incr n;
         Bytes.unsafe_set s' !n c
       | '\n' ->
         Bytes.unsafe_set s' !n '\\';
         incr n;
         Bytes.unsafe_set s' !n 'n'
       | '\t' ->
         Bytes.unsafe_set s' !n '\\';
         incr n;
         Bytes.unsafe_set s' !n 't'
       | '\r' ->
         Bytes.unsafe_set s' !n '\\';
         incr n;
         Bytes.unsafe_set s' !n 'r'
       | '\b' ->
         Bytes.unsafe_set s' !n '\\';
         incr n;
         Bytes.unsafe_set s' !n 'b'
       | '%' when i + 1 < len && s.[i + 1] = '{' ->
         Bytes.unsafe_set s' !n '\\';
         incr n;
         Bytes.unsafe_set s' !n '%'
       | ' ' .. '~' as c -> Bytes.unsafe_set s' !n c
       | c ->
         let a = Char.code c in
         Bytes.unsafe_set s' !n '\\';
         incr n;
         Bytes.unsafe_set s' !n (Char.unsafe_chr (48 + (a / 100)));
         incr n;
         Bytes.unsafe_set s' !n (Char.unsafe_chr (48 + (a / 10 mod 10)));
         incr n;
         Bytes.unsafe_set s' !n (Char.unsafe_chr (48 + (a mod 10))));
      incr n
    done
  ;;

  (* Surround [s] with quotes, escaping it if necessary. *)
  let quoted s =
    let len = String.length s in
    let n = quote_length s in
    let s' = Bytes.create (n + 2) in
    Bytes.unsafe_set s' 0 '"';
    if len = 0 || n > len
    then escape_to s ~dst:s' ~ofs:1
    else Bytes.blit_string ~src:s ~src_pos:0 ~dst:s' ~dst_pos:1 ~len;
    Bytes.unsafe_set s' (n + 1) '"';
    Bytes.unsafe_to_string s'
  ;;
end

type t =
  | Quoted_string of string
  | List of t list

let rec to_string t =
  match t with
  | Quoted_string s -> Escape.quoted s
  | List l -> Printf.sprintf "(%s)" (List.map l ~f:to_string |> String.concat ~sep:" ")
;;
