open Format
open Common.Pp

type 'a t = { desc : 'a desc; meta : 'a }

and 'a desc =
  | Const of Const.t
  | Var of string
  | Lam of string list * 'a t
  | App of 'a t * 'a t list
  | Prim1 of Prim.op1 * 'a t
  | Prim2 of Prim.op2 * 'a t * 'a t
  | If of 'a t * 'a t * 'a t
  | Let of 'a binding * 'a t
  | Letrec of 'a binding list * 'a t
  | Tuple of 'a t list
  | Get of 'a t * int
  | Vect of 'a t * 'a t
  | VectGet of 'a t * 'a t
  | VectSet of 'a t * 'a t * 'a t
  | VectLength of 'a t
  | Ascribe of 'a t * Type.t
  | For of string * 'a t * 'a t * 'a t
  | Variant of string * 'a t list * Type.t
  | Match of 'a t * 'a pattern list * Type.t
  | Seq of 'a t list

and 'a binding = { name : string; term : 'a t }
and 'a pattern = { casename : string; args : string list; body : 'a t }

let rec pp ?(show_meta = false) pp_meta fmt e =
  let pp_with_meta fmt e =
    if show_meta then
      fprintf fmt "@[<hv>%a@ %a@]" pp_meta e.meta (pp_desc pp_meta show_meta) e
    else pp_desc pp_meta show_meta fmt e
  in
  pp_with_meta fmt e

and pp_desc pp_meta show_meta fmt e =
  let pp = pp ~show_meta pp_meta in
  match e.desc with
  | Const const -> fprintf fmt "%s" (Const.to_string const)
  | Var x -> fprintf fmt "%s" x
  | Lam (args, e) ->
      fprintf fmt "@[(fun %a ->@;<1 2>%a)@]"
        (pp_list ~sep:" " pp_print_string)
        args pp e
  | App (e1, exprs) -> pp_app pp fmt (e1 :: exprs)
  | Prim1 (op1, e) -> fprintf fmt "%a %a" Prim.pp_op1 op1 pp e
  | Prim2 (op2, e1, e2) ->
      fprintf fmt "@[(%a %a@;<1 2>%a)@]" pp e1 Prim.pp_op2 op2 pp e2
  | If (e1, e2, e3) ->
      fprintf fmt "@[<hv>(if %a then@;<1 2>%a@;<1 1>else@;<1 2>%a)@]" pp e1 pp
        e2 pp e3
  | Let (binding, e) ->
      fprintf fmt "@[<v>@[<b>let@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]"
        (pp_binding ~show_meta pp_meta)
        binding pp e
  | Letrec (bindings, e) ->
      fprintf fmt "@[<v>@[<b>let rec@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]"
        (pp_list ~sep:"@;<1 0>and " (pp_binding ~show_meta pp_meta))
        bindings pp e
  | Tuple exprs -> fprintf fmt "@[(%a)@]" (pp_list ~sep:",@;<1 2>" pp) exprs
  | Get (e, i) -> fprintf fmt "%a.(%d)" pp e i
  | Vect (e1, e2) -> pp_app_name "vector" pp fmt [ e1; e2 ]
  | VectGet (e1, e2) -> pp_app_name "vector_get" pp fmt [ e1; e2 ]
  | VectSet (e1, e2, e3) -> pp_app_name "vector_set" pp fmt [ e1; e2; e3 ]
  | VectLength e -> pp_app_name "vector_length" pp fmt [ e ]
  | Ascribe (e', t) ->
      fprintf fmt "@[<hv>(%a@;<0 1>%a@;<1 1>:: %a)@]" pp_meta e.meta pp e'
        Type.pp t
  | For (x, e1, e2, e3) ->
      fprintf fmt "@[<v>@[for %s = %a to %a do@]@;<1 2>%a@;<1 0>done@]" x pp e1
        pp e2 pp e3
  | Variant (name, exprs, _) ->
      if List.length exprs = 0 then fprintf fmt "%s" name
      else fprintf fmt "@[%s (@[%a@])@]" name (pp_list ~sep:",@ " pp) exprs
  | Match (e, patterns, _) ->
      fprintf fmt "@[<hv>@[match %a@ with@]@;<1 0>| %a@]" pp e
        (pp_list ~sep:"@;<1 0>| " (pp_pattern ~show_meta pp_meta))
        patterns
  | Seq exprs ->
      fprintf fmt "@[<v>(%a)@]"
        (pp_list ~sep:";@;<1 1>" (fun fmt e -> fprintf fmt "@[%a@]" pp e))
        exprs

and pp_binding ?(show_meta = false) pp_meta fmt binding =
  fprintf fmt "%s =@;<1 2>%a" binding.name (pp ~show_meta pp_meta) binding.term

and pp_pattern ?(show_meta = false) pp_meta fmt pattern =
  fprintf fmt "@[%s %a ->@ %a@]" pattern.casename
    (pp_list ~sep:" " pp_print_string)
    pattern.args (pp ~show_meta pp_meta) pattern.body
