open Bwd

type t = Data.syn =
  | Var of int
  | Pi of t * (* BINDS *) t
  | Lam of (* BINDS *) t
  | App of t * t
  | Sg of t * (* BINDS *) t
  | Pair of t * t
  | Fst of t
  | Snd of t
  | Type
  | Bool
  | True
  | False
  | BoolElim of {
      motive : (* BINDS *) t;
      true_case : t;
      false_case : t;
      scrut : t;
    }

module Pretty =
struct
  type local_names = Yuujinchou.Trie.path option bwd
  type tm_with_names = t * local_names
  
  let pp_name fmt (id : Yuujinchou.Trie.path) = Format.pp_print_list
      ~pp_sep:(fun fmt () -> Format.pp_print_char fmt '.')
      Format.pp_print_string fmt id

  module Internal =
  struct
    type env = {
      prec : int;
    }

    module Eff = Algaeff.Reader.Make (struct type nonrec t = env end)

    let with_prec (prec : int) (inner : 'a Fmt.t) (fmt : Format.formatter) (v : 'a) =
      Eff.scope (fun _ -> { prec; }) (fun () -> inner fmt v)

    let delimited inner = with_prec 0 inner

    let parenthesize (prec : int) : 'a Fmt.t -> 'a Fmt.t =
      if (Eff.read()).prec > prec then
        Fmt.parens
      else
        Fun.id

    let lassoc (prec : int)
        ?(pre = "")
        (pp_left : 'l Fmt.t)
        ?(middle = " ")
        (pp_right : 'r Fmt.t)
        ?(post = "") : ('l * 'r) Fmt.t = 
      parenthesize prec @@ fun fmt v -> 
      Fmt.pf fmt "%s@[%a@]%s@[%a@]%s" 
        pre
        (with_prec prec pp_left) (fst v)
        middle 
        (with_prec (prec + 1) pp_right) (snd v)
        post

    let rassoc (prec : int)
        ?(pre = "")
        (pp_left : 'l Fmt.t)
        ?(middle = " ")
        (pp_right : 'r Fmt.t)
        ?(post = "") : ('l * 'r) Fmt.t = 
      parenthesize prec @@ fun fmt v -> 
      Fmt.pf fmt "%s@[%a@]%s@[%a@]%s" 
        pre
        (with_prec (prec + 1) pp_left) (fst v)
        middle 
        (with_prec prec pp_right) (snd v)
        post

    let pp_var ix fmt = Fmt.pf fmt "%@%d" ix

    let rec pp_pair fmt = (delimited @@ fun fmt (a, b) -> Fmt.pf fmt "@[(@[%a@], @[%a@])@]" pp a pp b) fmt

    and pp_bool_elim fmt =
      (delimited @@ fun fmt (motive, true_case, false_case, scrut) -> Fmt.pf fmt
         "@[bool-elim @[%a@] at _ -> @[%a@] [ true => @[%a@], false => @[%a@] ]@]"
         pp scrut pp motive pp true_case pp false_case
      ) fmt

    (* TODO: improve indentation *)
    and pp fmt = function
      | Var ix -> pp_var ix fmt
      | Pi (base, fam) -> rassoc 1 ~pre:"(_ : " pp ~middle:") -> " pp fmt (base, fam)
      | Lam body -> Fmt.pf fmt "@[\\_ -> @[%a@]@]" pp body
      | Sg (base, fam) -> rassoc 2 ~pre:"(_ : " pp ~middle:") * " pp fmt (base, fam)
      | App (a, b) -> lassoc 3 pp pp fmt (a, b)
      | Pair (a, b) -> pp_pair fmt (a, b)
      | Fst a -> lassoc 3 Fmt.string pp fmt ("fst", a)
      | Snd a -> lassoc 3 Fmt.string pp fmt ("snd", a)
      | Type -> Fmt.string fmt "type"
      | Bool -> Fmt.string fmt "bool"
      | True -> Fmt.string fmt "true"
      | False -> Fmt.string fmt "false"
      | BoolElim { motive; true_case; false_case; scrut } -> pp_bool_elim fmt (motive, true_case, false_case, scrut)
  end

  let pp fmt tm =
    let env : Internal.env = { prec = 0; } in
    Internal.Eff.run ~env (fun () -> Internal.pp fmt tm)
end