exception Json_error of string

type json =
  | JObject of (string * json) list
  | JArray of json list
  | JString of string
  | JBool of bool
  | JNull

let read_file path =
  let ic = open_in_bin path in
  Fun.protect
    ~finally:(fun () -> close_in_noerr ic)
    (fun () ->
      let n = in_channel_length ic in
      really_input_string ic n)

let parse_json s =
  let n = String.length s in
  let i = ref 0 in
  let error m = raise (Json_error (m ^ " at byte " ^ string_of_int !i)) in
  let ws () =
    while !i < n && match s.[!i] with ' ' | '\n' | '\r' | '\t' -> true | _ -> false do
      incr i
    done
  in
  let expect c =
    ws ();
    if !i >= n || s.[!i] <> c then error ("expected " ^ String.make 1 c);
    incr i
  in
  let parse_lit lit v =
    let len = String.length lit in
    if !i + len <= n && String.sub s !i len = lit then (i := !i + len; v)
    else error ("expected " ^ lit)
  in
  let parse_string () =
    expect '"';
    let b = Buffer.create 16 in
    let rec loop () =
      if !i >= n then error "unterminated string";
      match s.[!i] with
      | '"' -> incr i; JString (Buffer.contents b)
      | '\\' ->
          incr i;
          if !i >= n then error "unterminated escape";
          let c =
            match s.[!i] with
            | '"' -> '"'
            | '\\' -> '\\'
            | '/' -> '/'
            | 'b' -> '\b'
            | 'f' -> '\012'
            | 'n' -> '\n'
            | 'r' -> '\r'
            | 't' -> '\t'
            | _ -> error "unsupported string escape"
          in
          Buffer.add_char b c;
          incr i;
          loop ()
      | c ->
          Buffer.add_char b c;
          incr i;
          loop ()
    in
    loop ()
  in
  let rec value () =
    ws ();
    if !i >= n then error "expected value";
    match s.[!i] with
    | '{' -> obj ()
    | '[' -> arr ()
    | '"' -> parse_string ()
    | 't' -> parse_lit "true" (JBool true)
    | 'f' -> parse_lit "false" (JBool false)
    | 'n' -> parse_lit "null" JNull
    | _ -> error "unexpected value"
  and obj () =
    expect '{';
    ws ();
    if !i < n && s.[!i] = '}' then (incr i; JObject [])
    else
      let rec fields acc =
        let key =
          match parse_string () with JString k -> k | _ -> assert false
        in
        expect ':';
        let v = value () in
        ws ();
        if !i < n && s.[!i] = ',' then (incr i; fields ((key, v) :: acc))
        else if !i < n && s.[!i] = '}' then (incr i; JObject (List.rev ((key, v) :: acc)))
        else error "expected ',' or '}'"
      in
      fields []
  and arr () =
    expect '[';
    ws ();
    if !i < n && s.[!i] = ']' then (incr i; JArray [])
    else
      let rec elems acc =
        let v = value () in
        ws ();
        if !i < n && s.[!i] = ',' then (incr i; elems (v :: acc))
        else if !i < n && s.[!i] = ']' then (incr i; JArray (List.rev (v :: acc)))
        else error "expected ',' or ']'"
      in
      elems []
  in
  let v = value () in
  ws ();
  if !i <> n then error "trailing input";
  v

let member k = function
  | JObject xs -> begin
      match List.assoc_opt k xs with
      | Some v -> v
      | None -> raise (Json_error ("missing field: " ^ k))
    end
  | _ -> raise (Json_error "expected object")

let to_string = function
  | JString s -> s
  | _ -> raise (Json_error "expected string")

let to_list = function
  | JArray xs -> xs
  | _ -> raise (Json_error "expected array")

let to_assoc = function
  | JObject xs -> xs
  | _ -> raise (Json_error "expected object")

let rec named_of_json j =
  match member "tag" j |> to_string with
  | "Var" -> Named.Var (member "name" j |> to_string)
  | "Lam" ->
      Named.Lam (member "param" j |> to_string, named_of_json (member "body" j))
  | "App" ->
      Named.App (named_of_json (member "fun" j), named_of_json (member "arg" j))
  | "Subst" ->
      Named.Subst
        ( member "var" j |> to_string,
          named_of_json (member "repl" j),
          named_of_json (member "body" j) )
  | t -> raise (Json_error ("unknown tag: " ^ t))

let rec json_of_named = function
  | Named.Var x -> JObject [ ("tag", JString "Var"); ("name", JString x) ]
  | Named.Lam (x, b) ->
      JObject [ ("tag", JString "Lam"); ("param", JString x); ("body", json_of_named b) ]
  | Named.App (f, a) ->
      JObject [ ("tag", JString "App"); ("fun", json_of_named f); ("arg", json_of_named a) ]
  | Named.Subst (x, s, b) ->
      JObject
        [ ("tag", JString "Subst");
          ("var", JString x);
          ("repl", json_of_named s);
          ("body", json_of_named b) ]

type corpus_entry = {
  name : string;
  t1 : Named.t;
  t2 : Named.t;
  expected : bool option;
}

let load_corpus path =
  parse_json (read_file path) |> to_list
  |> List.map (fun e ->
         {
           name = member "name" e |> to_string;
           t1 = named_of_json (member "t1" e);
           t2 = named_of_json (member "t2" e);
           expected =
             (match member "expected" e with
             | JNull -> None
             | JBool b -> Some b
             | _ -> raise (Json_error "expected must be bool or null"));
         })

let load_bench path =
  parse_json (read_file path) |> to_assoc
  |> List.map (fun (k, v) -> (k, named_of_json v))