type t = Atom of string | List of t list

let parse text =
  let length = String.length text in
  let position = ref 0 in
  let fail message = invalid_arg (Printf.sprintf "%s at byte %d" message !position) in
  let rec space () =
    if !position < length then
      match text.[!position] with
      | ' ' | '\n' | '\r' | '\t' -> incr position; space ()
      | _ -> ()
  in
  let quoted () =
    incr position;
    let buffer = Buffer.create 32 in
    let rec loop () =
      if !position = length then fail "unterminated string";
      let character = text.[!position] in
      incr position;
      match character with
      | '"' -> Buffer.contents buffer
      | '\\' ->
          if !position = length then fail "incomplete escape";
          let escaped = match text.[!position] with
            | 'n' -> '\n' | 'r' -> '\r' | 't' -> '\t'
            | '"' -> '"' | '\\' -> '\\'
            | _ -> fail "unsupported escape"
          in
          incr position;
          Buffer.add_char buffer escaped;
          loop ()
      | character -> Buffer.add_char buffer character; loop ()
    in
    loop ()
  in
  let rec value () =
    space ();
    if !position = length then fail "expected expression";
    match text.[!position] with
    | '(' ->
        incr position;
        let rec items result =
          space ();
          if !position = length then fail "unterminated list";
          if text.[!position] = ')' then begin
            incr position;
            List (List.rev result)
          end else items (value () :: result)
        in
        items []
    | ')' -> fail "unexpected closing parenthesis"
    | '"' -> Atom (quoted ())
    | _ ->
        let start = !position in
        while !position < length && not (List.mem text.[!position] ['('; ')'; ' '; '\n'; '\r'; '\t']) do
          incr position
        done;
        Atom (String.sub text start (!position - start))
  in
  let rec values result =
    space ();
    if !position = length then List.rev result else values (value () :: result)
  in
  values []

let one text = match parse text with
  | [value] -> value
  | _ -> invalid_arg "expected one expression"

let rec print = function
  | Atom text ->
      if text <> "" && not (String.exists (fun c -> List.mem c ['('; ')'; ' '; '\n'; '\r'; '\t'; '"'; '\\']) text)
      then text else begin
        let buffer = Buffer.create (String.length text + 2) in
        Buffer.add_char buffer '"';
        String.iter (function
          | '"' -> Buffer.add_string buffer "\\\""
          | '\\' -> Buffer.add_string buffer "\\\\"
          | '\n' -> Buffer.add_string buffer "\\n"
          | '\r' -> Buffer.add_string buffer "\\r"
          | '\t' -> Buffer.add_string buffer "\\t"
          | c -> Buffer.add_char buffer c) text;
        Buffer.add_char buffer '"';
        Buffer.contents buffer
      end
  | List values -> "(" ^ String.concat " " (List.map print values) ^ ")"

let atom = function Atom text -> text | _ -> invalid_arg "expected atom"
let list = function List values -> values | _ -> invalid_arg "expected list"
let field name values =
  let matches = List.filter_map (function
    | List (Atom key :: values) when key = name -> Some values
    | _ -> None) values
  in
  match matches with
  | [values] -> values
  | [] -> invalid_arg ("missing field " ^ name)
  | _ -> invalid_arg ("duplicate field " ^ name)

let integer value = int_of_string (atom value)
let word = function
  | Atom text when String.starts_with ~prefix:"#x" text ->
      Int64.of_string ("0x" ^ String.sub text 2 (String.length text - 2))
  | Atom text when String.starts_with ~prefix:"#b" text ->
      Int64.of_string ("0b" ^ String.sub text 2 (String.length text - 2))
  | Atom text -> Int64.of_string text
  | List [Atom "-"; Atom text] -> Int64.neg (Int64.of_string text)
  | _ -> invalid_arg "invalid scalar"
let a text = Atom text
let i value = Atom (string_of_int value)
let w value = Atom (Printf.sprintf "0x%016Lx" value)