open Bwd
open Bwd.Infix

module S = NbE.Syntax

module Internal =
struct
  type env = {
    names : Name.local bwd;
    prec : int;
  }

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

  let bind (name : Name.local) f fmt x =
    let update env = { env with names = env.names <: name } in
    Eff.scope update (fun () -> f fmt x)

  let with_prec (prec : int) (inner : 'a Fmt.t) (fmt : Format.formatter) (v : 'a) =
    Eff.scope (fun env -> { env with 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)
      (pp_left : 'l Fmt.t)
      ?(sep = " ")
      (pp_right : 'r Fmt.t) : ('l * 'r) Fmt.t = 
    parenthesize prec @@ fun fmt v -> 
    Fmt.pf fmt "@[%a@]%s@[%a@]" 
      (with_prec prec pp_left) (fst v)
      sep
      (with_prec (prec + 1) pp_right) (snd v)

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

  let pp_local_name fmt name =
    Fmt.string fmt @@ Option.value ~default:"_" name

  let pp_var fmt ix =
    let names = (Eff.read()).names in
    match Bwd.nth_opt names ix with
    (* the variable has a bound name which has not been shadowed *)
    | Some ((Some name) as id) when Bwd.find_index ((=) id) names = Some ix -> Fmt.string fmt name
    (* the variable does not have a bound name or its bound name has been shadowed *)
    | Some _ -> Fmt.pf fmt "%@%d" ix
    | None -> Error.bug "index out of range in pp_var"

  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_var, motive, true_case, false_case, scrut) -> Fmt.pf fmt
       "@[bool-elim @[%a@] at @[%a@] -> @[%a@] [ true => @[%a@], false => @[%a@] ]@]"
       pp scrut pp_local_name motive_var (bind motive_var pp) motive pp true_case pp false_case
    ) fmt

  and pp_arg fmt (name, tp) = match name with
    | Some name -> (delimited @@ fun fmt (name, tp) -> Fmt.pf fmt "(@[%s@] : @[%a@])" name pp tp) fmt (name, tp)
    | None -> pp fmt tp

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

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