(* Copyright 2012 Tony Garnock-Jones <tonygarnockjones@gmail.com>. *)

(* This file is part of Hop. *)

(* Hop is free software: you can redistribute it and/or modify it *)
(* under the terms of the GNU General Public License as published by the *)
(* Free Software Foundation, either version 3 of the License, or (at your *)
(* option) any later version. *)

(* Hop is distributed in the hope that it will be useful, but *)
(* WITHOUT ANY WARRANTY; without even the implied warranty of *)
(* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU *)
(* General Public License for more details. *)

(* You should have received a copy of the GNU General Public License *)
(* along with Hop.  If not, see <http://www.gnu.org/licenses/>. *)

(* SPKI SEXP *)

open Lwt
open Lwt_io

exception Syntax_error of string

type display_hint_t = {hint : bytes; body : bytes}
and t =
  | Str of bytes
  | Hint of display_hint_t
  | Arr of t list

let emptystr = Str Bytes.empty
let litstr s = Str (Bytes.unsafe_of_string s)
let str s = Str (Bytes.of_string s)

let compare a b = Stdlib.compare a b

let digit_val c = (int_of_char c) - (int_of_char '0')
let val_digit n = char_of_int (n + 48)

let bytes_of_int =
  let siz = 40 in
  let buf = Bytes.make siz (* enough for 128 bits *) ' ' in
  function
    | 0 -> Bytes.of_string "0"
    | n ->
      let rec loop n i =
	if n = 0
	then Bytes.sub buf (siz - i) i
	else (Bytes.unsafe_set buf (siz - i - 1) (val_digit (n mod 10));
	      loop (n / 10) (i + 1))
      in loop n 0

let int_of_bytes bs = int_of_string (Bytes.to_string bs)

let write_simple_string write s =
  lwt () = write (bytes_of_int (Bytes.length s)) in
  lwt () = write (Bytes.of_string ":") in
  write s

let generic_output_sexp write x =
  let writestr = write_simple_string write in
  let rec walk x =
    match x with
      | Str s ->
	writestr s
      | Hint {hint = h; body = b} ->
	lwt () = write (Bytes.of_string "[") in
	lwt () = writestr h in
	lwt () = write (Bytes.of_string "]") in
	writestr b
      | Arr xs ->
	lwt () = write (Bytes.of_string "(") in
	lwt () = Lwt_list.iter_s walk xs in
	write (Bytes.of_string ")")
  in walk x

let write_bytes ch bs = write_from_exactly ch bs 0 (Bytes.length bs)

let output_sexp ch x = generic_output_sexp (write_bytes ch) x
let stream_of_sexp x = Streamutil.stream_generator (fun yield -> generic_output_sexp yield x)

let write_char_escaped ch c =
  if c = '\"'
  then write ch "\\\""
  else write_char ch c

let write_simple_string_human ch s =
  lwt () = write_char ch '\"' in
  lwt () = write_bytes ch (Bytes.escaped s) in
  write_char ch '\"'

let rec output_sexp_human ch x =
  match x with
    | Str s ->
      write_simple_string_human ch s
    | Hint {hint = h; body = b} ->
      lwt () = write_char ch '[' in
      lwt () = write_simple_string_human ch h in
      lwt () = write_char ch ']' in
      write_simple_string_human ch b
    | Arr xs ->
      lwt () = write_char ch '(' in
      lwt () =
	(match xs with
	  | [] -> return ()
	  | [x] -> output_sexp_human ch x
	  | (x :: xs') ->
	    lwt () = output_sexp_human ch x in
	    Lwt_list.iter_s
	      (fun x ->
		lwt () = write_char ch ' ' in
		output_sexp_human ch x)
	      xs') in
      write_char ch ')'

let char_numeric c = '0' <= c && c <= '9'
let char_whitespace c = c <= ' '

let input_bytes ch count =
  let buf = Bytes.create count in
  lwt () = read_into_exactly ch buf 0 count in
  return buf

let syntax_error explanation = raise_lwt (Syntax_error explanation)

let input_sexp_outer input_char input_bytes =
  let rec input_simple_string len =
    match_lwt input_char () with
      | ':' -> lwt bs = input_bytes len in return bs
      | b when char_numeric b -> input_simple_string (len * 10 + digit_val b)
      | _ -> syntax_error "Bad simple-string length character"
  in
  let rec input_sexp_list () =
    let rec collect acc =
      match_lwt input_sexp_inner () with
	| None -> return (Arr (List.rev acc))
	| Some v -> collect (v :: acc)
    in collect []
  and input_sexp_inner () =
    match_lwt input_char () with
      | '(' -> lwt xs = input_sexp_list () in return (Some xs)
      | ')' -> return None
      | '[' ->
	lwt hint = input_simple_string 0 in
	(match_lwt input_char () with
	  | ']' -> lwt b = input_simple_string 0 in return (Some (Hint {hint = hint; body = b}))
	  | _ -> syntax_error "Missing close-bracket in display hint")
      | b when char_numeric b ->
	lwt s = input_simple_string (digit_val b) in return (Some (Str s))
      | b when char_whitespace b ->
	(* Convenience for testing *)
	input_sexp_inner ()
      | _ ->
	syntax_error "Bad SEXP input character"
  in
  match_lwt input_sexp_inner () with
    | None -> syntax_error "Unexpected end of list"
    | Some v -> return v

let input_sexp ch = input_sexp_outer (fun () -> read_char ch) (input_bytes ch)
let parse b =
  input_sexp_outer
    (fun () -> return (Ibuffer.next_char b))
    (fun count -> return (Ibuffer.next_chars b count))
let sexp_of_bytes s = parse (Ibuffer.of_bytes s)
let bytes_of_sexp x = Streamutil.stream_to_bytes (stream_of_sexp x)

let assoc' key v =
  match v with
  | Arr entries ->
      let rec search entries =
	match entries with
	| [] -> None
	| (Arr (k :: result :: _)) :: _ when k = key -> Some result
	| _ :: rest -> search rest
      in search entries
  | _ -> None

let assoc k default v =
  match assoc' k v with
  | Some result -> result
  | None -> default