(** Base types and functor to create XML documents. *)

(** {2 Locations} *)

(** A position in a source (file, string, ...). *)
type pos = Lexing.position

let pos ?(fname="") ~line ~bol ~char () =
  { Lexing.pos_fname = fname; pos_lnum = line ;
    pos_bol = bol ; pos_cnum = char ;
  }

(** A location is a range defined by two positions.*)
type loc = pos * pos
type 'a with_loc = 'a * loc
type 'a with_loc_option = 'a * loc option

(** [loc loc_start loc_stop] creates a {!type-loc} structure with
  the given positions. *)
let loc start stop = (start, stop)

(** [loc_of_pos p len] creates a {!type-loc} structure with
  starting at position [p] and ending at position [p + len].
  We aware that only the [pos_cnum] field is modified in
  the ending position.
 *)
let loc_of_pos pos len =
  loc pos
    { pos with Lexing.pos_cnum = pos.Lexing.pos_cnum + len }

let string_of_loc loc =
  let open Lexing in
  let (start, stop) = loc in
  let line = start.pos_lnum in
  let char = start.pos_cnum - start.pos_bol in
  let len =
    if start.pos_fname = stop.pos_fname then
      stop.pos_cnum - start.pos_cnum
    else
      1
  in
  let file = start.pos_fname in
  Printf.sprintf "%sline %d, character%s %d%s"
    (match file with
     | "" -> ""
     | _ -> Printf.sprintf "File %S, " file)
    line
    (if len > 1 then "s" else "")
    char
    (if len > 1 then Printf.sprintf "-%d" (char + len) else "")

let string_of_loc_option = function
  | None -> ""
  | Some loc -> string_of_loc loc

let pp_loc ppf loc = Format.pp_print_string ppf (string_of_loc loc)
let pp_loc_option ppf loc = Format.pp_print_string ppf (string_of_loc_option loc)

let loc_sprintf loc fmt =
  match loc with
  | None -> Printf.sprintf fmt
  | Some loc -> Printf.ksprintf
      (fun s -> Printf.sprintf "%s:\n%s" (string_of_loc loc) s)
      fmt

(** {2 Errors} *)

type error = ..
exception Error of error
let error e = raise (Error e)

(**/**)

let string_of_error_funs = ref []

(**/**)

let register_string_of_error f =
  string_of_error_funs := f :: !string_of_error_funs
let string_of_error =
  let rec iter e = function
  | [] -> "No printer for error"
  | f :: q ->
      match f e with
      | Some s -> s
      | None -> iter e q
  in
  fun e -> iter e !string_of_error_funs

let () = Printexc.register_printer
  (function
       | Error e -> Some (string_of_error e)
       | _ -> None)


(** {2 Documents} *)

type cdata = { loc: loc option; text: string; quoted: bool }
type comment = { loc: loc option; comment: string }

let compare_cdata d1 d2 = String.compare d1.text d2.text
let compare_comment c1 c2 = String.compare c1.comment c2.comment

module type P =
  sig
    module Attributes : Map.S
    type attr_value
    type data
    val compare_name : Attributes.key -> Attributes.key -> int
    val compare_attr_value : attr_value -> attr_value -> int
    val compare_data : data -> data -> int
    val default_data : unit -> data
    val version_name : unit ->  Attributes.key
    val default_version : unit -> attr_value
    val default_attr_value : unit -> attr_value
    val pp_name : Format.formatter -> Attributes.key -> unit

    (** Print attribute value, with escaping but without quotes around. *)
    val pp_attr_value : Format.formatter -> attr_value -> unit

    (** If provided, use this function to print attributes.
      Else {!pp_name} and {!pp_attr_value} will be used. *)
    val pp_attributes : (Format.formatter -> attr_value Attributes.t -> unit) option
  end
module type S =
  sig
    type name
    type attr_value
    type attributes
    type data

    type proc_inst = { loc: loc option; app: name; args: string}
    type xml_decl = { loc: loc option; atts: attributes }
    type doctype = { loc: loc option; name: name; args: string}

    type node =
     { loc: loc option;
       name: name ;
       atts: attributes ;
       subs: tree list ;
       data: data ;
     }
    and tree =
    | E of node
    | D of cdata
    | C of comment
    | PI of proc_inst

    type prolog_misc = PC of comment | PPI of proc_inst

    type  prolog = {
      decl : xml_decl option ;
      misc : prolog_misc list ;
      doctype : doctype option ;
    }

    type doc = { prolog : prolog ; elements : tree list }

    val compare_doc : doc -> doc -> int
    val compare_tree : tree -> tree -> int

    (** {2 Constructors} *)

    val node_ : ?loc:loc -> name -> ?atts:attributes ->
      ?data: data -> tree list -> node
    val node :  ?loc:loc -> name -> ?atts:attributes ->
      ?data: data -> tree list -> tree
    val cdata_ : ?loc:loc -> ?quoted:bool -> string -> cdata
    val cdata : ?loc:loc -> ?quoted:bool -> string -> tree
    val comment_ : ?loc:loc -> string -> comment
    val comment : ?loc:loc -> string -> tree
    val prolog_comment : ?loc:loc -> string -> comment
    val pi_ : ?loc:loc -> name -> string -> proc_inst
    val pi : ?loc:loc -> name -> string -> tree
    val prolog_pi : ?loc:loc -> name -> string -> proc_inst
    val xml_decl : ?loc:loc -> attributes -> xml_decl
    val doctype : ?loc:loc -> name -> string -> doctype
    val prolog : ?decl:xml_decl -> ?doctype:doctype ->
      prolog_misc list -> prolog
    val doc : prolog -> tree list -> doc

    val doc_empty : unit -> doc

    (** {2 Handling attributes} *)

    (** [get_att atts name] returns the value associated to given attribute, if any. *)
    val get_att : attributes -> name -> attr_value option

    (** [opt_att atts ?def name] returns the value associated to given
      attribute, or the default value specified by [?def]. If [def] is not
      specified, [P.default_attr_value] is used. *)
    val opt_att : attributes -> ?def:attr_value -> name -> attr_value

    (** [atts_of_list list] returns a {!type-attributes} structure from the given
      list of pairs [(name, value) ; (name, value) ; ...].
      @param atts can be used to specify an existing {!attributes} structure to
      add bindings to, instead of starting from an empty one. *)
    val atts_of_list : ?atts:attributes -> (name * attr_value) list -> attributes

    (** [atts_one ?atts name value] is like {!atts_of_list} but for one attribute. *)
    val atts_one : ?atts:attributes -> name -> attr_value -> attributes

    (** [atts_remove name attributes] removes the binding to [name] from the
      [attributes]. *)
    val atts_remove : name -> attributes -> attributes

    (** [atts_replace name value attributes] adds a new bindings from [name] to
       [value] in [attributes]. If [name] was previously bound, the previous
       binding is removed. *)
    val atts_replace : name -> attr_value -> attributes -> attributes

    val atts_fold : (name -> attr_value -> 'acc -> 'acc) -> attributes -> 'acc -> 'acc
    val atts_empty : unit -> attributes

    (** {2 Printers} *)

    val pp_doc : Format.formatter -> doc -> unit
    val pp : Format.formatter -> doc -> unit
    val pp_attr : ?first:bool -> Format.formatter -> name -> attr_value -> unit
    val pp_attributes : Format.formatter -> attributes -> unit
    val pp_trees : Format.formatter -> tree list -> unit
    val doc_to_string : doc -> string
    val to_string : tree list -> string
    val string_of_xml : tree -> string

    (** {2 Text extraction} *)

    val doctype_name : doc -> name option
    val text_of_xmls : tree list -> string
    val text_of_range : tree -> ?inc_start:bool -> tree ->
      ?inc_stop:bool -> tree option -> string
  end
module Make (P:P) : S with
  type name = P.Attributes.key and
  type attr_value = P.attr_value and
  type attributes = P.attr_value P.Attributes.t and
  type data = P.data =
  struct
    module Attributes = P.Attributes
    type name = Attributes.key
    type attr_value = P.attr_value
    type data = P.data
    type attributes = attr_value Attributes.t
    let atts_empty : unit -> attributes = fun () -> Attributes.empty

    type proc_inst = { loc: loc option; app: name; args: string}
    type xml_decl = { loc: loc option; atts: attributes }
    type doctype = { loc: loc option; name: name; args: string}

    type node =
     { loc: loc option;
       name: name ;
       atts: attributes ;
       subs: tree list ;
       data: data ;
     }
    and tree =
    | E of node
    | D of cdata
    | C of comment
    | PI of proc_inst

    type prolog_misc = PC of comment | PPI of proc_inst

    type prolog = {
      decl : xml_decl option ;
      misc : prolog_misc list ;
      doctype : doctype option ;
    }
    type doc = { prolog : prolog; elements : tree list }

    let compare_attributes = Attributes.compare P.compare_attr_value

    let compare_proc_inst p1 p2 =
        match P.compare_name p1.app p2.app with
        | n when n <> 0 -> n
        | _ -> String.compare p1.args p2.args

    let rec compare_node n1 n2 =
        match P.compare_name n1.name n2.name with
        | n when n <> 0 -> n
        | _ ->
            match compare_attributes n1.atts n2.atts with
            | n when n <> 0 -> n
            | _ ->
                match P.compare_data n1.data n2.data with
                | n when n <> 0 -> n
                | _ -> List.compare compare_tree n1.subs n2.subs

    and compare_tree t1 t2 =
        match t1, t2 with
        | E n1, E n2 -> compare_node n1 n2
        | E _, _ -> 1
        | _, E _ -> -1
        | D d1, D d2 -> compare_cdata d1 d2
        | D _, _ -> 1
        | _, D _ -> -1
        | C c1, C c2 -> compare_comment c1 c2
        | C _, _ -> 1
        | _, C _ -> -1
        | PI p1, PI p2 -> compare_proc_inst p1 p2

    let compare_xml_decl (d1:xml_decl) (d2:xml_decl) =
        compare_attributes d1.atts d2.atts

    let compare_prolog_misc m1 m2 =
        match m1, m2 with
        | PC c1, PC c2 -> compare_comment c1 c2
        | PC _, _ -> 1
        | _, PC _ -> -1
        | PPI p1, PPI p2 -> compare_proc_inst p1 p2

    let compare_doctype (t1:doctype) (t2:doctype) =
        match P.compare_name t1.name t2.name with
        | n when n <> 0 -> n
        | _ -> String.compare t1.args t2.args

    let compare_prolog p1 p2 =
        match Option.compare compare_xml_decl p1.decl p2.decl with
        | n when n <> 0 -> n
        | _ ->
            match List.compare compare_prolog_misc p1.misc p2.misc with
            | n when n <> 0 -> n
            | _ -> Option.compare compare_doctype p1.doctype p2.doctype

    let compare_doc d1 d2 =
        match compare_prolog d1.prolog d2.prolog with
        | n when n <> 0 -> n
        | _ -> List.compare compare_tree d1.elements d2.elements

    let node_ ?loc name ?(atts=atts_empty ()) ?(data=P.default_data()) subs =
      { loc; name; atts ; subs; data }
    let node ?loc name ?atts ?data subs = E (node_ ?loc name ?atts ?data subs)

    let cdata_ ?loc ?(quoted=false) text = { loc ; text ; quoted }
    let cdata ?loc ?quoted text = D (cdata_ ?loc ?quoted text)

    let comment_ ?loc comment = { loc ; comment }
    let comment ?loc c = C (comment_ ?loc c)

    let prolog_comment ?loc comment = { loc ; comment }

    let pi_ ?loc app args = { loc ; app ; args }
    let pi ?loc app args = PI (pi_ ?loc app args)

    let prolog_pi ?loc app args = { loc ; app ; args }
    let xml_decl ?loc atts : xml_decl = { loc ; atts }
    let doctype ?loc name args : doctype = { loc ; name ; args }
    let prolog ?decl ?doctype misc = { decl ; doctype ; misc }
    let doc prolog elements = { prolog ; elements }

    let doc_empty () = doc (prolog []) []

    let doctype_name doc = match doc.prolog.doctype with None -> None | Some dt -> Some dt.name

    let get_att atts name = Attributes.find_opt name atts

    let opt_att atts ?(def=P.default_attr_value()) name =
      match get_att atts name with None -> def | Some v -> v

    let atts_of_list =
       let f acc (name,v) = Attributes.add name v acc in
       fun ?(atts=Attributes.empty) l -> List.fold_left f atts l

    let atts_one ?(atts=Attributes.empty) name v = Attributes.add name v atts
    let atts_remove = Attributes.remove
    let atts_replace = Attributes.add
    let atts_fold = Attributes.fold

    (** {2 Printing} *)

    let pp_attr ?(first=false) ppf name value =
      Format.fprintf ppf "%s%a=\"%a\""
          (if first then "" else " ")
          P.pp_name name
          P.pp_attr_value value

    let pp_attributes ppf atts =
        match P.pp_attributes with
        | Some p -> p ppf atts
        | None ->
            match Attributes.bindings atts with
            | [] -> ()
            | (name,value)::q ->
                pp_attr ~first:true ppf name value;
                List.iter (fun (n,v) -> pp_attr ppf n v) q

    let rec pp_tree ppf = function
    | C { comment } ->
        Format.fprintf ppf "<!--%s-->" (Misc.escape comment)
    | D { text ; quoted = true } ->
        Format.fprintf ppf "<![CDATA[%s]]>" text
    | D { text ; quoted = false } ->
        Format.fprintf ppf "%s" (Misc.escape text)
    | PI { app ; args } ->
        Format.fprintf ppf "<?%a %s?>"
          P.pp_name app (Misc.escape args)
    | E { name ; atts ; subs } ->
        Format.fprintf ppf "<%a" P.pp_name name;
          if not (Attributes.is_empty atts) then
            Format.fprintf ppf " %a" pp_attributes atts;
        match subs with
          [] -> Format.pp_print_string ppf "/>"
        | _ ->
            Format.pp_print_string ppf ">";
            List.iter (pp_tree ppf) subs;
            Format.fprintf ppf "</%a>" P.pp_name name

    let pp_trees ppf l = List.iter (pp_tree ppf) l

    let pp_prolog_misc ppf = function
    | PC { comment } ->
        Format.fprintf ppf "<!--%s-->\n" (Misc.escape comment)
    | PPI { app ; args } ->
        Format.fprintf ppf "<?%a %s?>\n"
          P.pp_name app (Misc.escape args)

    let pp_prolog ppf p =
      (match p.decl with
       | None -> ()
       | Some { atts } ->
           Format.pp_print_string ppf "<?xml " ;
           let vname = P.version_name () in
           let v = opt_att ~def:(P.default_version()) atts vname in
           pp_attributes ppf (Attributes.singleton vname v);
           let atts = Attributes.remove vname atts in
           if not (Attributes.is_empty atts) then
               Format.fprintf ppf " %a" pp_attributes atts;
           Format.pp_print_string ppf "?>\n"
      );
      List.iter (pp_prolog_misc ppf) p.misc;
      (
       match p.doctype with
       | None -> ()
       | Some { name ; args } ->
           Format.fprintf ppf "<!DOCTYPE %a %s>\n"
             P.pp_name name (Misc.escape args)
      )

    let pp_doc ppf doc =
      pp_prolog ppf doc.prolog ;
      List.iter (pp_tree ppf) doc.elements

    let pp ppf doc = pp_trees ppf doc.elements

    let doc_to_string = Misc.string_of_pp pp_doc
    let to_string = Misc.string_of_pp pp_trees
    let string_of_xml = Misc.string_of_pp pp_tree
    let string_of_atts atts = Misc.string_of_pp pp_attributes atts

    let text_of_xmls =
      let rec iter acc = function
      | [] -> String.concat "" (List.rev acc)
      | E { subs } :: q -> iter acc (subs @ q)
      | D { text } :: q -> iter (text :: acc) q
      | _ :: q -> iter acc q
      in
      fun t -> iter [] t

    let text_of_range =
      (* `In_range true: include elements;
         `In_range false: in range but ignore first one 
      *)
      let rec iter b ~inc_start ~start ~inc_stop ~stop in_range t =
        match t, in_range, stop with
        | _, `After,  _-> `After
        | _, `Before, _ when t == start ->
            Log.debug (fun m -> m "Xtmpl.Types.text_of_range: start found");
            iter b ~inc_start ~start ~inc_stop ~stop (`In_range inc_start) t
        | D cdata, `In_range false, None -> `In_range true
        | D cdata, `In_range true, None -> Buffer.add_string b cdata.text; in_range
        | D cdata, `In_range inc, Some stop ->
            if t == stop then
              (if inc_stop then Buffer.add_string b cdata.text;
               `After
              )
            else
              (
               if inc then Buffer.add_string b cdata.text;
               `In_range true
              )
        | E { subs }, _, None -> 
            List.fold_left
              (iter b ~inc_start ~start ~inc_stop ~stop) 
              in_range subs
        | E { subs }, _, Some stop_ ->
            let in_range = List.fold_left
              (iter b ~inc_start ~start ~inc_stop ~stop) in_range subs
            in
            if t == stop_ then `After else in_range
        | D _, `Before, _ -> in_range
        (* beware that this make inc_stop only works when stop is a cdata *)
        | _, `In_range _, Some stop when t == stop -> `After
        | C _, _, _ -> in_range
        | PI _, _, _ -> in_range
      in
      fun root ?(inc_start=true) start ?(inc_stop=false) stop ->
        let b = Buffer.create 256 in
        let in_range = iter b ~inc_start ~start ~inc_stop ~stop `Before root in
        Log.debug (fun m -> m "Xtmpl.Types.text_of_range => %s"
           (match in_range with
            | `Before -> "`Before"
            | `After -> "`After"
            | `In_range b -> Printf.sprintf "`In_range (%b)" b));
        Buffer.contents b
  end

(** [merge_cdata c1 c2] returns a new {!type:cdata} by merging text
  of the ones provided. If both have a location, the returned
  [cdata] has a location, ranging from the start position of the
  first to the end position of the second. If any of the cdatas has
  its [quoted] flag set to [true], then the new cdata has this flag
  set to [true].*)
let merge_cdata (c1 : cdata) (c2 : cdata) =
  let loc =
    match c1.loc, c2.loc with
      None, _ | _, None -> None
    | Some l1, Some l2 -> Some (fst l1, snd l2)
  in
  { loc ; text = c1.text ^ c2.text ;
    quoted = c1.quoted || c2.quoted ;
  }

(** {2 Mapping trees} *)

(** Description of how to map various types in trees *)
type ('n1,'av1,'a1,'d1,'n2,'av2,'a2,'d2) map_param =
  { map_name : 'n1 -> 'n2 (** mapping names *);
    map_attr_value : 'av1 -> 'av2 (** mapping attribute values *);
    map_attributes : ('a1 -> 'a2) option
      (** mapping attributes; if provided attributes will be mapped by
          this function; else [map_name] and [map_attr_value] will be
         used for mapping each pair (name, value) of attributes. *) ;
    map_data : 'd1 -> 'd2 (** mapping data *) ;
  }

(** The functions to map trees and documents, obtained with
  {!Make_map.mapper}. *)
type ('a1,'x1, 'd1, 'p1, 'a2, 'x2, 'd2, 'p2) mapper = {
  map_atts: 'a1 -> 'a2 (** map attributes *) ;
  map_xml : 'x1 -> 'x2 (** map a single tree *) ;
  map_xmls : 'x1 list -> 'x2 list (** map a list of trees *) ;
  map_doc : 'd1 -> 'd2 (** map a whole document *) ;
  map_prolog : 'p1 -> 'p2 (** map prolog of a document *) ;
}

(** This functor creates a module with a {!Make_map.mapper} function
  from trees and documents of the first module to the second one. *)
module Make_map (X1:S) (X2:S) =
  struct
  (**/**)

  let map_proc_inst map (pi:X1.proc_inst) =
    X2.pi_ ?loc:pi.loc (map.map_name pi.app) pi.args
  let map_prolog_misc map = function
  | X1.PC c -> X2.PC c
  | PPI pi -> X2.PPI (map_proc_inst map pi)
  let rec map_xml map = function
    | X1.D d -> X2.D d
    | C c -> X2.comment ?loc:c.loc c.comment
    | PI pi -> X2.PI (map_proc_inst map pi)
    | E { loc ; name ; atts ; subs ; data } ->
        let atts = map_atts map atts in
        let subs = map_xmls map subs in
        let data = map.map_data data in
        X2.node ?loc (map.map_name name) ~data ~atts subs
  and map_xmls map l = List.map (map_xml map) l
  and map_atts map atts =
      match map.map_attributes with
      | Some f -> f atts
      | None ->
          X1.atts_fold
            (fun name value acc ->
               let name = map.map_name name in
               let value = map.map_attr_value value in
               X2.atts_one ~atts:acc name value
            )
            atts (X2.atts_empty())
  and map_prolog map (p:X1.prolog) =
      let decl =
        match p.decl with
        | None -> None
        | Some d ->
            let atts = map_atts map d.atts in
            Some(X2.xml_decl ?loc:d.loc atts)
      in
      let doctype = match p.doctype with
        | None -> None
        | Some d -> Some (X2.doctype ?loc:d.loc (map.map_name d.name) d.args)
      in
      let misc = List.map (map_prolog_misc map) p.misc in
      X2.prolog ?decl ?doctype misc

  and map_doc map (d:X1.doc) =
      let prolog = map_prolog map d.prolog in
      X2.doc prolog (map_xmls map d.elements)

  (**/**)

  let mapper = fun (map:(X1.name,X1.attr_value,X1.attributes,X1.data,
      X2.name,X2.attr_value,X2.attributes,X2.data) map_param) ->
      {
        map_xmls = map_xmls map ;
        map_atts = map_atts map ;
        map_xml = map_xml map ;
        map_doc = map_doc map ;
        map_prolog = map_prolog map ;
      }
end
