type t =
  | Cons   of t * t
  | Sym    of string
  | String of string
  | Int    of int
  | Float  of float

let nil = Sym "nil"

let escaped str =
  let len = String.length str in
  let extra_chars = ref 0 in
  for i = 0 to len - 1 do
    match str.[i] with
    | '\\' | '"' -> incr extra_chars
    | _ -> ()
  done;
  let buf = Buffer.create (len + !extra_chars + 2) in
  Buffer.add_char buf '"';
  if !extra_chars = 0 then (
    Buffer.add_string buf str
  ) else (
    for i = 0 to len - 1 do
      let c = str.[i] in
      if c = '"' || c = '\\' then
        Buffer.add_char buf '\\';
      Buffer.add_char buf c
    done;
  );
  Buffer.add_char buf '"';
  Buffer.contents buf

let unescaped str =
  (* Unescaped doesn't support unicode escaping and multibyte hex and octal
     escaping.
     Unicode escaping: '\uNNNN' or '\U00NNNNNN'
     Hex/octal escaping looks like '\xNN' or '\NNN'.
     '\xNNNN' and '\NNNNNN' are ambiguous, but emacs will try to parse them
     as multibyte
  *)
  match String.index str '\\' with
  | exception Not_found -> str
  | _ ->
    let len = String.length str in
    let buf = Buffer.create len in
    let i = ref 0 in
    while !i < len do
      match str.[!i] with
      | '\\' -> (
          incr i;
          begin match str.[!i] with
            | 'n' -> Buffer.add_char buf '\n'
            | 'r' -> Buffer.add_char buf '\r'
            | 't' -> Buffer.add_char buf '\t'
            | 'x' ->
              let c0 = Char.code str.[!i+1] in
              let c1 = Char.code str.[!i+2] in
              Buffer.add_char buf (Char.chr ((c0 * 16) lor c1));
              i := !i + 2;
            | '0'..'9' ->
              let c0 = Char.code str.[!i+1] in
              let c1 = Char.code str.[!i+2] in
              let c2 = Char.code str.[!i+3] in
              Buffer.add_char buf (Char.chr ((c0 * 64) lor (c1 * 8) lor c2));
              i := !i + 2;
            | c -> Buffer.add_char buf c
          end;
          incr i
        )
      | c ->
        Buffer.add_char buf c;
        incr i
    done;
    Buffer.contents buf

let rec of_list = function
  | [] -> nil
  | a :: tl -> Cons (a, of_list tl)

let rec tell_sexp tell = function
  | Cons (a,b) ->
    tell "(";
    tell_sexp tell a;
    tell_cons tell b
  | Sym s    -> tell s
  | String s -> tell (escaped s)
  | Int i    -> tell (string_of_int i)
  | Float f  -> tell (string_of_float f)

and tell_cons tell = function
  | Sym "nil" -> tell ")"
  | Cons (a,b) ->
    tell " ";
    tell_sexp tell a;
    tell_cons tell b
  | sexp ->
    tell " . ";
    tell_sexp tell sexp;
    tell ")"

let is_alpha c =
  (c >= 'a' && c <= 'z')
  || (c >= 'A' && c <= 'Z')

let is_num c =
  (c >= '0' && c <= '9' || c == '-')

let is_alphanum c = is_alpha c || is_num c

let read_sexp getch =
  let buf = Buffer.create 10 in
  let rec read_sexp getch = function
    | ' ' | '\t' | '\n' ->
      read_sexp getch (getch ())

    | c when is_num c ->
      read_num getch c

    | '\'' | ':' | '_' as c -> read_sym getch (Some c)
    | c when is_alpha c -> read_sym getch (Some c)

    | '"' ->
      read_string getch
    | '\000' -> raise End_of_file
    | '(' ->
      let lhs, next = read_sexp getch (getch ()) in
      read_cons getch (fun rhs -> Cons (lhs, rhs)) next
    | _ -> failwith "Invalid parse"

  and read_cons getch k next =
    match (match next with Some c -> c | None -> getch ()) with
    | ' ' | '\t' | '\n' -> read_cons getch k None
    | ')' -> k nil, None
    | '.' ->
      let rhs, next = read_sexp getch (getch ()) in
      let rec aux = function
        | ')' -> k rhs
        | ' ' | '\t' | '\n' -> aux (getch ())
        | _ -> failwith "Invalid parse"
      in
      begin match next with
        | Some c -> aux c
        | None -> aux (getch ())
      end, None
    | c ->
      let cell, next = read_sexp getch c in
      read_cons getch (fun rhs -> k (Cons (cell, rhs))) next

  and read_num getch c =
    Buffer.clear buf;
    Buffer.add_char buf c;
    let rec aux ~is_start ~is_float =
      match getch () with
      | '-' when is_start ->
        Buffer.add_char buf c; aux ~is_start:false ~is_float
      | c when c >= '0' && c <= '9' ->
        Buffer.add_char buf c; aux ~is_start:false ~is_float
      | '.' | 'e' | 'E' as c ->
        Buffer.add_char buf c; aux ~is_start:false ~is_float:true
      | c ->
        let s = Buffer.contents buf in
        (if is_float
         then Float (float_of_string s)
         else Int (int_of_string s)),
        Some c
    in
    aux ~is_start:true ~is_float:false

  and read_string getch =
    Buffer.clear buf;
    let rec aux () =
      match getch () with
      | '\000' -> failwith "Unterminated string"
      | '\\' ->
        Buffer.add_char buf '\\';
        Buffer.add_char buf (getch ());
        aux ()
      | '"' ->
        String (unescaped (Buffer.contents buf)), None
      | c ->
        Buffer.add_char buf c;
        aux ()
    in
    aux ()

  and read_sym getch next =
    Buffer.clear buf;
    let rec aux next =
      match (match next with Some c -> c | None -> getch ()) with
      | ('\'' | '-' | ':' | '_') as c ->
        Buffer.add_char buf c;
        aux None
      | c when is_alphanum c ->
        Buffer.add_char buf c;
        aux None
      | c -> Sym (Buffer.contents buf), Some c
    in
    aux next
  in
  read_sexp getch (getch ())

let to_buf sexp buf =
  tell_sexp (Buffer.add_string buf) sexp

let to_string sexp =
  let buf = Buffer.create 100 in
  to_buf sexp buf;
  Buffer.contents buf

let getch_of_substring str pos len =
  let len = pos + len in
  if pos < 0 || len > String.length str then
    invalid_arg "Sexp.getch_of_substring";
  let pos = ref pos in
  let getch () =
    if !pos < len then
      let r = str.[!pos] in
      incr pos;
      r
    else '\000'
  in
  getch

let getch_of_string str =
  getch_of_substring str 0 (String.length str)

let of_string str =
  fst (read_sexp (getch_of_string str))

let getch_of_subbytes str pos len =
  let len = pos + len in
  if pos < 0 || len > Bytes.length str then
    invalid_arg "Sexp.getch_of_subbytes";
  let pos = ref pos in
  let getch () =
    if !pos < len then
      let r = Bytes.get str !pos in
      incr pos;
      r
    else '\000'
  in
  getch

let of_file_descr ?(on_read=ignore) fd =
  let getch = ref (fun () -> '\000') in
  let rest = ref None in
  let buffer = Bytes.create 1024 in
  let getch () =
    match !rest with
    | Some r ->
      rest := None;
      r
    | None ->
      match !getch () with
      | '\000' ->
        on_read fd;
        let read = Unix.read fd buffer 0 1024 in
        if read = 0 then '\000'
        else
          begin
            getch := getch_of_subbytes buffer 0 read;
            !getch ()
          end
      | c -> c
  in
  fun () ->
    try
      let sexp, rest' = read_sexp getch in
      rest := rest';
      Some sexp
    with End_of_file -> None

let of_channel ?on_read ic =
  of_file_descr ?on_read (Unix.descr_of_in_channel ic)

let rec of_json =
  let assoc_item (a,b) = Cons (Sym a, of_json b) in
  function
  | `Null       -> Sym "null"
  | `Int i      -> Int i
  | `Float f    -> Float f
  | `String s   -> String s
  | `Bool true  -> Sym "true"
  | `Bool false -> Sym "false"
  | `Assoc lst  -> Cons (Cons (Sym "assoc", Sym "nil"), of_list (List.map assoc_item lst))
  | `List lst   -> of_list (List.map of_json lst)

let rec to_json =
  let fail msg sexp =
    failwith (msg ^ ", got: \n" ^ to_string sexp)
  in
  let rec assoc_item = function
    | Cons (Cons (Sym a, b), c) -> (a, to_json b) :: assoc_item c
    | Sym "nil" -> []
    | sexp -> fail "expecting association (key . value)" sexp
  in
  let rec list_items = function
    | Sym "nil" -> []
    | Cons (hd, tl) -> to_json hd :: list_items tl
    | sexp -> fail "expecting list" sexp
  in
  function
  | Sym "null"  -> `Null
  | Sym "true"  -> `Bool true
  | Sym "false" -> `Bool false
  | Int i    -> `Int i
  | Float f  -> `Float f
  | String s -> `String s
  | Cons (Cons (Sym "assoc", Sym "nil"), assocs) ->
    `Assoc (assoc_item assocs)
  | Sym "nil" -> `List []
  | Cons (hd, tl) -> `List (to_json hd :: list_items tl)
  | Sym s -> `String s
