type t =
  | Integer of Z.t
  | Rational of Q.t
  | String of string
  | Reference of string
  | Nil
  | Record of t Label.Map.t

let integer z = Integer z
let string s = String s
let nil = Nil

let reference name =
  let is_letter c = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') in
  let is_body c = is_letter c || (c >= '0' && c <= '9') || c = '_' in
  if
    String.length name > 0
    && is_letter name.[0]
    && String.for_all is_body name
  then Ok (Reference name)
  else Error (`Invalid_reference name)

let rational ~num ~den =
  if Z.equal den Z.zero then Error `Zero_denominator
  else
    (* Q.make reduces by the gcd and moves the sign to the
       numerator; collapsing den = 1 to Integer is the minimal-shape
       rule that makes integer inclusion literal. *)
    let q = Q.make num den in
    if Z.equal (Q.den q) Z.one then Ok (Integer (Q.num q))
    else Ok (Rational q)

let record bindings =
  match bindings with
  | [] -> Error `Empty_record
  | _ ->
      let step acc (label, term) =
        match acc with
        | Error _ as e -> e
        | Ok map ->
            if Label.Map.mem label map then Error (`Duplicate_label label)
            else Ok (Label.Map.add label term map)
      in
      Result.map
        (fun map -> Record map)
        (List.fold_left step (Ok Label.Map.empty) bindings)

(* The head/tail labels of sequence cells. Direct construction:
   both spellings are well formed by inspection. *)
let head_label =
  match Label.of_string "head" with Ok l -> l | Error _ -> assert false

let tail_label =
  match Label.of_string "tail" with Ok l -> l | Error _ -> assert false

let cons ~head ~tail =
  Record Label.Map.(empty |> add head_label head |> add tail_label tail)

let sequence items = List.fold_right (fun x acc -> cons ~head:x ~tail:acc) items Nil

let rec compare a b =
  match (a, b) with
  | Integer x, Integer y -> Z.compare x y
  | Rational x, Rational y -> Q.compare x y
  | String x, String y -> String.compare x y
  | Reference x, Reference y -> String.compare x y
  | Nil, Nil -> 0
  | Record x, Record y -> Label.Map.compare compare x y
  | Integer _, _ -> -1
  | _, Integer _ -> 1
  | Rational _, _ -> -1
  | _, Rational _ -> 1
  | String _, _ -> -1
  | _, String _ -> 1
  | Reference _, _ -> -1
  | _, Reference _ -> 1
  | Nil, _ -> -1
  | _, Nil -> 1

let equal a b = compare a b = 0

(* A record prints as #( ... ) exactly when it is a cell of the
   sequence elaboration: the two labels head and tail, nothing else. *)
let as_cell = function
  | Record map when Label.Map.cardinal map = 2 -> (
      match
        (Label.Map.find_opt head_label map, Label.Map.find_opt tail_label map)
      with
      | Some head, Some tail -> Some (head, tail)
      | _ -> None)
  | _ -> None

let rec pp fmt t =
  match t with
  | Integer z -> Z.pp_print fmt z
  | Rational q -> Q.pp_print fmt q
  | String s -> Format.fprintf fmt "%S" s
  | Reference name -> Format.fprintf fmt "#%s" name
  | Nil -> Format.pp_print_string fmt "#()"
  | Record map -> (
      match as_cell t with
      | Some (head, tail) -> pp_sequence fmt head tail
      | None ->
          Format.fprintf fmt "@[<hv 1>{";
          let first = ref true in
          Label.Map.iter
            (fun label term ->
              if not !first then Format.fprintf fmt "@ ";
              first := false;
              Format.fprintf fmt "%a %a" Label.pp label pp term)
            map;
          Format.fprintf fmt "}@]")

and pp_sequence fmt head tail =
  Format.fprintf fmt "@[<hv 2>#(%a" pp head;
  let rec spine = function
    | Nil -> ()
    | t -> (
        match as_cell t with
        | Some (h, rest) ->
            Format.fprintf fmt "@ %a" pp h;
            spine rest
        | None ->
            (* An improper chain: honest about the dangling tail. *)
            Format.fprintf fmt "@ | %a" pp t)
  in
  spine tail;
  Format.fprintf fmt ")@]"