(*********************************************************************************)
(*                OCaml-RDF                                                      *)
(*                                                                               *)
(*    Copyright (C) 2012-2024 Institut National de Recherche en Informatique     *)
(*    et en Automatique. All rights reserved.                                    *)
(*                                                                               *)
(*    This program is free software; you can redistribute it and/or modify       *)
(*    it under the terms of the GNU Lesser General Public License version        *)
(*    3 as published by the Free Software Foundation.                            *)
(*                                                                               *)
(*    This program is distributed in the hope that it will be useful,            *)
(*    but WITHOUT ANY WARRANTY; without even the implied warranty of             *)
(*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *)
(*    GNU General Public License for more details.                               *)
(*                                                                               *)
(*    You should have received a copy of the GNU General Public License          *)
(*    along with this program; if not, write to the Free Software                *)
(*    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA                   *)
(*    02111-1307  USA                                                            *)
(*                                                                               *)
(*    Contact: Maxence.Guesdon@inria.fr                                          *)
(*                                                                               *)
(*********************************************************************************)

(** *)

open Term;;

type options = (string * string) list
let get_option ?def name l =
  try List.assoc name l
  with Not_found ->
    match def with
        None -> failwith (Printf.sprintf "Missing option %S" name)
      | Some v -> v
;;

(** Interface to query Basic Graph Patterns (BGP) in a graph. *)
module type Storage_BGP =
  sig
    type g
    type term
    val term : g -> Term.term -> term
    val compare : g -> term -> term -> int
    val rdfterm : g -> term -> Term.term
    val subjects : g -> term list
    val objects : g -> term list
    val find :
        ?sub:term ->
        ?pred: term ->
        ?obj:term -> g -> (term * term * term) list
  end
;;

module type Storage =
  sig
    val name : string

    type g

    type error
    exception Error of error
    val string_of_error : error -> string

    val open_graph : ?options: (string * string) list -> Iri.t -> g
    val graph_name : g -> Iri.t
    val graph_size : g -> int

    val add_triple : g -> sub: term -> pred: Iri.t -> obj: term -> unit
    val rem_triple : g -> sub: term -> pred: Iri.t -> obj: term -> unit

    val add_triple_t : g -> triple -> unit
    val rem_triple_t : g -> triple -> unit

    val subjects_of : g -> pred: Iri.t -> obj: term -> term list
    val predicates_of : g -> sub: term -> obj: term -> Iri.t list
    val objects_of : g -> sub: term -> pred: Iri.t -> term list

    val find : ?sub: term -> ?pred: Iri.t -> ?obj: term -> g -> triple list
    val exists : ?sub: term -> ?pred: Iri.t -> ?obj: term -> g -> bool
    val exists_t : triple -> g -> bool

    val subjects : g -> term list
    val predicates : g -> Iri.t list
    val objects : g -> term list

    val folder : g -> Term.TSet.t Iri.Map.t Term.TMap.t option

    val transaction_start : g -> unit
    val transaction_commit : g -> unit
    val transaction_rollback : g -> unit

    val copy : g -> g

    val new_blank_id : g -> Term.blank_id

    val namespaces : g -> (Iri.t * string) list
    val add_namespace : g -> Iri.t -> string -> unit
    val rem_namespace : g -> string -> unit
    val set_namespaces : g -> (Iri.t * string) list -> unit

    module BGP : Storage_BGP with type g = g
  end

exception Storage_error of string * string * exn
let () = Printexc.register_printer
  (function
   | Storage_error (st, err, _) -> Some (Printf.sprintf "Graph storage %S: %s" st err)
   | _ -> None)

module Make (S : Storage) =
  struct
    type g = S.g

    let embed f x =
      try f x
      with (S.Error e) as exn ->
        raise (Storage_error (S.name, S.string_of_error e, exn))

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

    let open_graph ?options name = embed (S.open_graph ?options) name
    let graph_name = embed S.graph_name
    let graph_size = embed S.graph_size

    let add_triple g ~sub ~pred ~obj = embed (fun g -> S.add_triple g ~sub ~pred ~obj) g
    let rem_triple g ~sub ~pred ~obj = embed (fun g -> S.rem_triple g ~sub ~pred ~obj) g

    let add_triple_t g = embed (S.add_triple_t g)
    let rem_triple_t g = embed (S.rem_triple_t g)

    let subjects_of g ~pred ~obj = embed (fun g -> S.subjects_of g ~pred ~obj) g
    let predicates_of g ~sub ~obj = embed (fun g -> S.predicates_of g ~sub ~obj) g
    let objects_of g ~sub ~pred = embed (fun g -> S.objects_of g ~sub ~pred) g

    let find ?sub ?pred ?obj = embed (S.find ?sub ?pred ?obj)
    let exists ?sub ?pred ?obj = embed (S.exists ?sub ?pred ?obj)
    let exists_t triple = embed (S.exists_t triple)

    let subjects = embed S.subjects
    let predicates = embed S.predicates
    let objects = embed S.objects

    let folder = embed S.folder

    let transaction_start = embed S.transaction_start
    let transaction_commit = embed S.transaction_commit
    let transaction_rollback = embed S.transaction_rollback

    let copy = embed S.copy

    let new_blank_id = embed S.new_blank_id

    let namespaces = embed S.namespaces
    let add_namespace = embed S.add_namespace
    let rem_namespace = embed S.rem_namespace
    let set_namespaces = embed S.set_namespaces
    module BGP = S.BGP
  end

module type Graph =
  sig
    type g

    val open_graph : ?options: (string * string) list -> Iri.t -> g
    val graph_name : g -> Iri.t
    val graph_size : g -> int

    val add_triple : g -> sub: term -> pred: Iri.t -> obj: term -> unit
    val rem_triple : g -> sub: term -> pred: Iri.t -> obj: term -> unit

    val add_triple_t : g -> triple -> unit
    val rem_triple_t : g -> triple -> unit

    val subjects_of : g -> pred: Iri.t -> obj: term -> term list
    val predicates_of : g -> sub: term -> obj: term -> Iri.t list
    val objects_of : g -> sub: term -> pred: Iri.t -> term list

    val find : ?sub: term -> ?pred: Iri.t -> ?obj: term -> g -> triple list
    val exists : ?sub: term -> ?pred: Iri.t -> ?obj: term -> g -> bool
    val exists_t : triple -> g -> bool

    val subjects : g -> term list
    val predicates : g -> Iri.t list
    val objects : g -> term list

    val folder : g -> Term.TSet.t Iri.Map.t Term.TMap.t option
    val transaction_start : g -> unit
    val transaction_commit : g -> unit
    val transaction_rollback : g -> unit

    val copy : g -> g

    val new_blank_id : g -> Term.blank_id

    val namespaces : g -> (Iri.t * string) list
    val add_namespace : g -> Iri.t -> string -> unit
    val rem_namespace : g -> string -> unit
    val set_namespaces : g -> (Iri.t * string) list -> unit

    module BGP : Storage_BGP with type g = g
  end

let storages = ref [];;
let add_storage m =
  let module P = (val m : Storage) in
  let module M = Make (P) in
  storages := (P.name, (module M : Graph)) :: !storages
;;

type graph =
  {
    name : unit -> Iri.t ;
    size : unit -> int ;
    add_triple : sub: term -> pred: Iri.t -> obj: term -> unit ;
    rem_triple : sub: term -> pred: Iri.t -> obj: term -> unit ;
    add_triple_t : triple -> unit ;
    rem_triple_t : triple -> unit ;
    subjects_of : pred: Iri.t -> obj: term -> term list ;
    predicates_of : sub: term -> obj: term -> Iri.t list ;
    objects_of : sub: term -> pred: Iri.t -> term list ;
    find : ?sub: term -> ?pred: Iri.t -> ?obj: term -> unit -> triple list ;
    exists : ?sub: term -> ?pred: Iri.t -> ?obj: term -> unit -> bool ;
    exists_t : triple -> bool ;
    subjects : unit -> term list ;
    predicates : unit -> Iri.t list ;
    objects : unit -> term list ;
    folder : unit -> Term.TSet.t Iri.Map.t Term.TMap.t option ;
    transaction_start : unit -> unit ;
    transaction_commit : unit -> unit ;
    transaction_rollback : unit -> unit ;
    copy : unit -> graph ;
    new_blank_id : unit -> Term.blank_id ;
    namespaces : unit -> (Iri.t * string) list ;
    add_namespace : Iri.t -> string -> unit ;
    rem_namespace : string -> unit ;
    set_namespaces : (Iri.t * string) list -> unit ;

    bgp : (module Bgp.S) ;
  }


let open_graph ?(options=[]) name =
  let kind = get_option ~def: "mem" "storage" options in
  let storage =
    try List.assoc kind !storages
    with Not_found -> failwith (Printf.sprintf "Unknown storage %S" kind)
  in
  let module S = (val storage) in
  let rec mk g =
    let module P = struct
      type term = S.BGP.term
      type g = S.g
      let term = S.BGP.term g
      let compare = S.BGP.compare g
      let rdfterm = S.BGP.rdfterm g
      let subjects () = S.BGP.subjects g
      let objects () = S.BGP.objects g
      let find ?sub ?pred ?obj () = S.BGP.find ?sub ?pred ?obj g
      end
    in
    let module BGP = Bgp.Make (P) in
    let g =
      { name = (fun () -> S.graph_name g) ;
        size = (fun () -> S.graph_size g) ;
        add_triple = S.add_triple g ;
        rem_triple = S.rem_triple g ;
        add_triple_t = S.add_triple_t g ;
        rem_triple_t = S.rem_triple_t g ;
        subjects_of = S.subjects_of g ;
        predicates_of = S.predicates_of g ;
        objects_of = S.objects_of g ;
        find = (fun ?sub ?pred ?obj () -> S.find ?sub ?pred ?obj g) ;
        exists = (fun ?sub ?pred ?obj () -> S.exists ?sub ?pred ?obj g) ;
        exists_t = (fun t -> S.exists_t t g) ;
        subjects = (fun () -> S.subjects g) ;
        predicates = (fun () -> S.predicates g) ;
        objects = (fun () -> S.objects g) ;
        folder = (fun () -> S.folder g) ;
        transaction_start = (fun () -> S.transaction_start g) ;
        transaction_commit = (fun () -> S.transaction_commit g) ;
        transaction_rollback = (fun () -> S.transaction_rollback g) ;
        copy = (fun () -> mk (S.copy g)) ;
        new_blank_id = (fun () -> S.new_blank_id g) ;
        namespaces = (fun () -> S.namespaces g) ;
        add_namespace = S.add_namespace g ;
        rem_namespace = S.rem_namespace g ;
        set_namespaces = S.set_namespaces g ;
        bgp = (module BGP : Bgp.S) ;
      }
    in
    g.add_namespace (Rdf_.rdf_"") "rdf";
    g
  in
  let g = S.open_graph ~options name in
  let g = mk g in
  g
;;

module Bid_map = Map.Make
  (struct
     type t = Term.blank_id
     let compare id1 id2 =
       String.compare
         (Term.string_of_blank_id id1)
         (Term.string_of_blank_id id2)
   end
  );;

let merge ?(map=fun _ -> None) g1 g2 =
  let fmap bid_map x =
    match map x with
      Some x -> (bid_map, x)
    | None ->
        match x with
          Term.Iri _
        | Term.Literal _
        | Blank -> (bid_map, x)
        | Blank_ id ->
            let (id2, bid_map) =
              try (Bid_map.find id bid_map, bid_map)
              with Not_found ->
                  let id2 = g1.new_blank_id () in
                  let bid_map = Bid_map.add id id2 bid_map in
                  (id2, bid_map)
            in
            (bid_map, Blank_ id2)
  in
  let f bid_map (sub,pred,obj) =
    let (bid_map, sub) = fmap bid_map sub in
    let (bid_map, _) = fmap bid_map (Term.Iri pred) in
    let (bid_map, obj) = fmap bid_map obj in
    g1.add_triple ~sub ~pred ~obj;
    bid_map
  in
  let triples = g2.find () in
  ignore(List.fold_left f Bid_map.empty triples)
;;

let only_iris =
  let f acc = function
  | Term.Iri iri -> iri :: acc
  | _ -> acc
  in
  fun l -> List.rev (List.fold_left f [] l)

let only_literals =
  let f acc = function
  | Term.Literal lit -> lit :: acc
  | _ -> acc
  in
  fun l -> List.rev (List.fold_left f [] l)

let iri_objects_of g ~sub ~pred =
  only_iris (g.objects_of ~sub ~pred)

let iri_subjects_of g ~pred ~obj =
  only_iris (g.subjects_of ~pred ~obj)

let literal_objects_of g ~sub ~pred =
  only_literals (g.objects_of ~sub ~pred)

let subgraph_from ?options g root =
  let root_iri =
    match root with
      Term.Iri iri -> iri
    | Term.Blank | Term.Blank_ _
    | Term.Literal _ -> Iri.of_string ""
  in
  let g2 = open_graph ?options root_iri in
  let module S = Term.TSet in
  let rec iter seen sub =
    if S.mem sub seen then
      seen
    else
      let seen = S.add sub seen in
      let triples = g.find ~sub () in
      List.fold_left iter_triple seen triples
  and iter_triple seen ((sub,pred,obj) as t) =
    g2.add_triple_t t;
    iter seen obj
  in
  ignore(iter S.empty root) ;
  g2

let to_list =
  let rec iter g acc sub =
    match g.objects_of ~sub ~pred:Rdf_.first with
    | [] -> acc
    | t :: _ ->
        let acc = t :: acc in
        match g.objects_of ~sub ~pred:Rdf_.rest with
        | q :: _ -> iter g acc q
        | _ -> acc
  in
  fun g t -> List.rev (iter g [] t)

let add_list =
  let rec iter g tail = function
  | [] -> tail
  | h :: q ->
      let n = Term.blank_ (g.new_blank_id ()) in
      g.add_triple ~sub:n ~pred:Rdf_.rest ~obj:tail;
      g.add_triple ~sub:n ~pred:Rdf_.first ~obj:h;
      iter g n q
  in
  fun g l -> iter g (Term.Iri Rdf_.nil) (List.rev l)

let types_of g sub = iri_objects_of g ~sub ~pred:Rdf_.type_

let root_opt =
  let rec iter g seen sub =
    if Term.TSet.mem sub seen then
      Some sub
    else
      match g.find ~obj:sub () with
      | [] -> Some sub
      | (s,_,_) :: _ -> iter g (Term.TSet.add sub seen) s
  in
  fun g ->
    match g.subjects () with
    | [] -> None
    | sub :: _ -> iter g Term.TSet.empty sub

let root g = match root_opt g with None -> Term.Iri (g.name ()) | Some t -> t

type diff =
| Cardinals of int * int
| Missing_triple of Term.triple
| Extra_triple of Term.triple

let string_of_diff = function
| Cardinals (c1, c2) -> Printf.sprintf "Cardinals differ: %d <> %d" c1 c2
| Missing_triple (sub,pred,obj) ->
    let details title = function
    | Term.Iri i -> Printf.sprintf "\n%s: %s" title (Iri.to_string_details i)
    | _ -> ""
    in
    Printf.sprintf "Missing triple %s %s %s%s%s"
      (Term.string_of_term sub) Term.(string_of_term (Iri pred)) (Term.string_of_term obj)
      (details "subject" sub) (details "object" obj)
| Extra_triple (sub,pred,obj) ->
    Printf.sprintf "Extra triple %s %s %s"
      (Term.string_of_term sub) Term.(string_of_term (Iri pred)) (Term.string_of_term obj)

let pp_diff ppf diff = Format.fprintf ppf "%s" (string_of_diff diff)

let isomorphic_diff =
  let f ~ignore_blanks f_diff g2 acc (sub, pred, obj) =
    match acc with
    | Some _ -> acc
    | None ->
        match sub, obj with
        | (Term.Blank | Term.Blank_ _), _
        | _, (Term.Blank | Term.Blank_ _) when ignore_blanks -> acc
        | _ ->
            match g2.find ~sub ~pred ~obj () with
            | [] -> Some (f_diff (sub, pred, obj))
            | _ -> acc
  in
  fun ?(ignore_blanks=false) g1 g2 ->
    match List.fold_left
      (f ~ignore_blanks (fun x -> Missing_triple x) g2) None (g1.find())
    with
    | Some d -> Some d
    | None ->
        match List.fold_left
          (f ~ignore_blanks (fun x -> Extra_triple x) g1) None (g2.find())
        with
        | None when ignore_blanks ->
            let c1 = g1.size () in
            let c2 = g2.size () in
            if c1 <> c2 then
              Some (Cardinals (c1, c2))
            else
              None
        | x -> x


let copy_from =
  let module S = Term.TSet in
  let rec f keep g1 g2 seen (sub, pred, obj) =
    g2.add_triple ~sub ~pred ~obj;
    match obj with
    | Term.Blank_ _ -> fold keep g1 g2 seen obj
    | Term.Iri iri when keep iri -> fold keep  g1 g2 seen obj
    | _ -> seen
  and fold keep g1 g2 seen sub =
    if S.mem sub seen then
      seen
    else
      let seen = S.add sub seen in
      List.fold_left (f keep g1 g2) seen (g1.find ~sub ())
  in
  fun ?(keep=(fun _ -> false)) g term ->
    let g2 = open_graph (g.name()) in
    ignore (fold keep g g2 S.empty term);
    g2

let replace_objs g ~sub ~pred objs =
  List.iter g.rem_triple_t (g.find ~sub ~pred ());
  List.iter (fun obj -> g.add_triple ~sub ~pred ~obj) objs
