(************************************************************************)
(*         *      The Rocq Prover / The Rocq Development Team           *)
(*  v      *         Copyright INRIA, CNRS and contributors             *)
(* <O___,, * (see version control and CREDITS file for authors & dates) *)
(*   \VV/  **************************************************************)
(*    //   *    This file is distributed under the terms of the         *)
(*         *     GNU Lesser General Public License Version 2.1          *)
(*         *     (see LICENSE file for the text of the license)         *)
(************************************************************************)

open Util
open Pp
open Genarg
open Locus

let beautify_comments = ref []

let extract_comments pos =
  (* Take all comments that terminates before pos, or begin exactly
     at pos (used to print comments attached after an expression) *)
  let is_before ((b,e),_) = e <= pos || Int.equal pos b in
  let before, after = List.partition is_before !beautify_comments in
  beautify_comments := after;
  List.rev_map snd before

let pr_located pr (loc, x) =
  match loc with
  | Some loc ->
    let (b, e) = Loc.unloc loc in
    (* Side-effect: order matters *)
    let before = Pp.comment (extract_comments b) in
    let x = pr x in
    let after = Pp.comment (extract_comments e) in
    before ++ x ++ after
  | _ -> pr x

let pr_ast pr { CAst.loc; v } = pr_located pr (loc, v)

let pr_lident lid = pr_ast Names.Id.print lid

let pr_lname lna = pr_ast Names.Name.print lna

let pr_or_var pr = function
  | ArgArg x -> pr x
  | ArgVar id -> pr_lident id

let pr_or_by_notation f = let open Constrexpr in CAst.with_val (function
    | AN v -> f v
    | ByNotation (s,sc) -> qs s ++ pr_opt (fun sc -> str "%" ++ str sc) sc)

let hov_if_not_empty n p = if Pp.ismt p then p else hov n p

let rec pr_raw_generic env sigma ?level (GenArg (Rawwit wit, x)) =
  match wit with
    | ListArg wit ->
      let map x = pr_raw_generic env sigma ?level (in_gen (rawwit wit) x) in
      let ans = pr_sequence map x in
      hov_if_not_empty 0 ans
    | OptArg wit ->
      let ans = match x with
        | None -> mt ()
        | Some x -> pr_raw_generic env sigma ?level (in_gen (rawwit wit) x)
      in
      hov_if_not_empty 0 ans
    | PairArg (wit1, wit2) ->
      let p, q = x in
      let p = in_gen (rawwit wit1) p in
      let q = in_gen (rawwit wit2) q in
      hov_if_not_empty 0 (pr_sequence (pr_raw_generic env sigma ?level) [p; q])
    | ExtraArg s ->
       let open Genprint in
       match generic_raw_print (in_gen (rawwit wit) x) with
       | PrinterBasic pp -> pp env sigma
       | PrinterNeedsLevel { default_already_surrounded; printer } ->
         let level = Option.default default_already_surrounded level in
         printer env sigma level


let rec pr_glb_generic env sigma ?level (GenArg (Glbwit wit, x)) =
  match wit with
    | ListArg wit ->
      let map x = pr_glb_generic env sigma ?level (in_gen (glbwit wit) x) in
      let ans = pr_sequence map x in
      hov_if_not_empty 0 ans
    | OptArg wit ->
      let ans = match x with
        | None -> mt ()
        | Some x -> pr_glb_generic env sigma ?level (in_gen (glbwit wit) x)
      in
      hov_if_not_empty 0 ans
    | PairArg (wit1, wit2) ->
      let p, q = x in
      let p = in_gen (glbwit wit1) p in
      let q = in_gen (glbwit wit2) q in
      let ans = pr_sequence (pr_glb_generic env sigma ?level) [p; q] in
      hov_if_not_empty 0 ans
    | ExtraArg s ->
       let open Genprint in
       match generic_glb_print (in_gen (glbwit wit) x) with
       | PrinterBasic pp -> pp env sigma
       | PrinterNeedsLevel { default_already_surrounded; printer } ->
         let level = Option.default default_already_surrounded level in
         printer env sigma level