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

let rec of_term (term : Term.t) =
  match term with
  | Term.Integer z -> Integer z
  | Term.Rational q -> Rational q
  | Term.String s -> String s
  | Term.Nil -> Nil
  | Term.Record map -> Record (Label.Map.map of_term map)

let rec to_term = function
  | Var _ -> None
  | Integer z -> Some (Term.integer z)
  | Rational q ->
      (* Rebuilding through the smart constructor keeps canonicity a
         single-point invariant. The denominator is nonzero by
         construction, so this cannot fail. *)
      Result.to_option (Term.rational ~num:(Q.num q) ~den:(Q.den q))
  | String s -> Some (Term.string s)
  | Nil -> Some Term.nil
  | Record map ->
      let closed =
        Label.Map.fold
          (fun label pattern acc ->
            match (acc, to_term pattern) with
            | Some bindings, Some term -> Some ((label, term) :: bindings)
            | _ -> None)
          map (Some [])
      in
      Option.bind closed (fun bindings -> Result.to_option (Term.record bindings))

let is_ground pattern = Option.is_some (to_term pattern)

let vars pattern =
  let rec collect acc = function
    | Var v -> if List.mem v acc then acc else v :: acc
    | Integer _ | Rational _ | String _ | Nil -> acc
    | Record map -> Label.Map.fold (fun _ p acc -> collect acc p) map acc
  in
  List.rev (collect [] pattern)

let forced_label name =
  match Label.of_string name with Ok l -> l | Error _ -> assert false

let head_label = forced_label "head"
let tail_label = forced_label "tail"

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

let rec equal a b =
  match (a, b) with
  | Var x, Var y -> String.equal x y
  | Integer x, Integer y -> Z.equal x y
  | Rational x, Rational y -> Q.equal x y
  | String x, String y -> String.equal x y
  | Nil, Nil -> true
  | Record x, Record y -> Label.Map.equal equal x y
  | _ -> false

let strictly_contains pattern ~sub =
  let rec inside p = equal p sub || strictly p
  and strictly = function
    | Record map -> Label.Map.exists (fun _ p -> inside p) map
    | _ -> false
  in
  strictly pattern

let rec pp fmt = function
  | Var v -> Format.pp_print_string fmt v
  | Integer z -> Z.pp_print fmt z
  | Rational q -> Q.pp_print fmt q
  | String s -> Format.fprintf fmt "%S" s
  | Nil -> Format.pp_print_string fmt "#()"
  | Record map ->
      Format.fprintf fmt "@[<hv 1>{";
      let first = ref true in
      Label.Map.iter
        (fun label p ->
          if not !first then Format.fprintf fmt "@ ";
          first := false;
          Format.fprintf fmt "%a %a" Label.pp label pp p)
        map;
      Format.fprintf fmt "}@]"

module Subst = struct
  module M = Map.Make (String)

  type pattern = t
  type nonrec t = pattern M.t

  let empty = M.empty

  let rec walk subst pattern =
    match pattern with
    | Var v -> (
        match M.find_opt v subst with
        | Some bound -> walk subst bound
        | None -> pattern)
    | _ -> pattern

  let rec resolve subst pattern =
    match walk subst pattern with
    | Record map -> Record (Label.Map.map (resolve subst) map)
    | other -> other

  let rec occurs subst v pattern =
    match walk subst pattern with
    | Var w -> String.equal v w
    | Record map -> Label.Map.exists (fun _ p -> occurs subst v p) map
    | _ -> false

  let rec unify subst a b =
    let a = walk subst a and b = walk subst b in
    match (a, b) with
    | Var x, Var y when String.equal x y -> Some subst
    | Var x, other | other, Var x ->
        if occurs subst x other then None else Some (M.add x other subst)
    | Integer x, Integer y -> if Z.equal x y then Some subst else None
    | Rational x, Rational y -> if Q.equal x y then Some subst else None
    | String x, String y -> if String.equal x y then Some subst else None
    | Nil, Nil -> Some subst
    | Record x, Record y ->
        if
          Label.Map.cardinal x = Label.Map.cardinal y
          && Label.Map.for_all (fun label _ -> Label.Map.mem label y) x
        then
          Label.Map.fold
            (fun label p acc ->
              Option.bind acc (fun s -> unify s p (Label.Map.find label y)))
            x (Some subst)
        else None
    | _ -> None
end