(**************************************************************************)
(*                                                                        *)
(*  Copyright (C) 2012 Johannes 'josch' Schauer <j.schauer@email.de>      *)
(*  Copyright (C) 2012 Pietro Abate <pietro.abate@pps.jussieu.fr>         *)
(*                                                                        *)
(*  This library is free software: you can redistribute it and/or modify  *)
(*  it under the terms of the GNU Lesser General Public License as        *)
(*  published by the Free Software Foundation, either version 3 of the    *)
(*  License, or (at your option) any later version.  A special linking    *)
(*  exception to the GNU Lesser General Public License applies to this    *)
(*  library, see the COPYING file for more information.                   *)
(**************************************************************************)

open ExtLib
open Common
open Debian
open Algo

include Util.Logging(struct let label = __FILE__ end) ;;

let progressbar_build_graph = Util.Progress.create "build_graph"
let progressbar_build_graph_bfs = Util.Progress.create "build_graph_bfs"
let timer_build_graph = Util.Timer.create "build_graph"
let timer_build_graph_bfs = Util.Timer.create "build_graph_bfs"

module IntSet = BootstrapCommon.IntSet

type vertex =
  |InstSet of (int * IntSet.t)
  |SrcPkg  of int

module Unique = Unique.Make(struct
  type t = vertex
  let compare v1 v2 =
    match v1,v2 with 
    |SrcPkg p1, SrcPkg p2 -> Pervasives.compare p1 p2
    |InstSet (p1,s1), InstSet (p2,s2) -> begin
        match Pervasives.compare p1 p2 with
          | 0 -> IntSet.compare s1 s2
          | i -> i
     end
    |SrcPkg _, InstSet _ -> -1
    |InstSet _, SrcPkg _ -> 1
end)

type kind =
  |BuildDep
  |BuildsFrom of IntSet.t ref

type 'a edge = { 
  depend : kind;
  annotation : 'a list
}

let default_edge = { depend = BuildsFrom (ref IntSet.empty); annotation = [] }

type et = [`Loop | `NoLoop | `StrongDep ]

let compare_edge e1 e2 = compare e1 e2

module PkgV = struct
  type t = Unique.t
  let compare v1 v2 = Pervasives.compare (Unique.uid v1) (Unique.uid v2)
  let hash = Unique.uid
  let equal v1 v2 = (Unique.uid v1) = (Unique.uid v2)
end
 
module PkgE = struct
  type t = et edge
  let compare x y = compare_edge x y
  let equal x y = ( compare_edge x y ) = 0
  let default = default_edge
end

module G = Graph.Imperative.Digraph.ConcreteBidirectionalLabeled(PkgV)(PkgE)

module VertexSet = Set.Make(G.V)
module EdgeSet = Set.Make(G.E)

(* given a list of tuples where the first element is a source package and the
 * second element is a partitioning of its installation set between its direct
 * dependencies, connect binary packages to source packages according to
 * "bin2src" and omit connections of binary packages to source packages for
 * binary packages that qualify as "available" *)
let graph_helper available bin2src univ instsets =
  let g = G.create () in
  List.iter (fun (srcpkg, partitions) ->
    let src = Unique.create (SrcPkg (CudfAdd.vartoint univ srcpkg)) in
    List.iter (fun (pid,is) ->
        let instset = Unique.create (InstSet (pid,is)) in
        let label = { depend = BuildDep; annotation = [] } in
        let edge = (src,label,instset) in
        G.add_edge_e g edge;
        IntSet.iter (fun pid ->
          let binpkg = CudfAdd.inttovar univ pid in
          if not (available binpkg) then begin
            let srcpkgdep = bin2src binpkg in
            let dst = Unique.create (SrcPkg (CudfAdd.vartoint univ srcpkgdep)) in
            try match G.E.label (G.find_edge g instset dst) with
              |{ depend = BuildsFrom s } -> s := IntSet.add pid !s
              |_ -> assert false 
            with Not_found ->  begin
              let label = { depend = BuildsFrom (ref (IntSet.singleton pid)); annotation = [] } in
              let edge = (instset,label,dst) in
              G.add_edge_e g edge
            end
          end
        ) is
    ) partitions;
  ) instsets;
  g
;;

let get_src_package ?(allowmismatch=false) univ pkg = try
    BootstrapCommon.get_src_package ~allowmismatch univ pkg
  with Sources.NotfoundSrc ->
    failwith (Printf.sprintf "can't find source package for binary package %s" (CudfAdd.string_of_package pkg))
;;

(* create and time a normal buildgraph by calculating installation sets *)
let dist_graph ?(global_constraints=true) ?(available=(fun _ -> true)) ?(allowmismatch=false) ?(opt=false) ?(num_cores=1) custom_is_ht univ sl =
  Util.Timer.start timer_build_graph;
  Util.Progress.set_total progressbar_build_graph (List.length sl);

  let bin2src = get_src_package ~allowmismatch univ in
  let pool = Depsolver_int.init_pool_univ ~global_constraints univ in

  let opt_partition_cache = Hashtbl.create 10000 in
  let partitionofsrc = if opt then
      BootstrapCommon.compute_dependency_sets_opt ~global_constraints ~partition:true ~available opt_partition_cache pool univ
    else
      BootstrapCommon.compute_dependency_sets ~global_constraints ~partition:true custom_is_ht pool univ
  in

  let instsets = Parmap.parmap ~ncores:num_cores (fun srcpkg ->
    Util.Progress.progress progressbar_build_graph;
    let _, partitions = partitionofsrc srcpkg in
    (srcpkg, partitions)
  ) (Parmap.L sl) in

  Util.Timer.stop timer_build_graph (graph_helper available bin2src univ instsets)
;;

(* create and time a buildgraph by calculating installation closures *)
let closure_graph ?(global_constraints=true) ?(available=(fun _ -> true)) ?(allowmismatch=false) ?(num_cores=1) univ sl =
  Util.Timer.start timer_build_graph;
  Util.Progress.set_total progressbar_build_graph (List.length sl);

  let bin2src = get_src_package ~allowmismatch univ in
  let pool = Depsolver_int.init_pool_univ ~global_constraints univ in
  let globalid = Cudf.universe_size univ in

  let to_set l = List.fold_right IntSet.add l IntSet.empty in
  (* rule: connect everything to everything else *)
  let partitionofsrc src =
    let l = List.fold_left (fun acc vpkglist ->
      let disj = CudfAdd.resolve_vpkgs_int univ vpkglist in
      List.fold_left (fun acc2 pid ->
          let dc = Depsolver_int.dependency_closure_cache pool [pid; globalid] in
          let dcs = IntSet.remove globalid (to_set dc) in
          (pid, dcs)::acc2
      ) acc disj
    ) [] src.Cudf.depends in
    IntSet.empty, l
  in

  let instsets = Parmap.parmap ~ncores:num_cores (fun srcpkg ->
    Util.Progress.progress progressbar_build_graph;
    let _, partitions = partitionofsrc srcpkg in
    (srcpkg, partitions)
    ) (Parmap.L sl) in

  Util.Timer.stop timer_build_graph (graph_helper available bin2src univ instsets)
;;

(* create and time a normal buildgraph by calculating strong dependencies *)
let strong_graph ?(global_constraints=true) ?(available=(fun _ -> true)) ?(allowmismatch=false) ?(num_cores=1) univ sl =
  Util.Timer.start timer_build_graph;
  Util.Progress.set_total progressbar_build_graph (List.length sl);

  let bin2src = get_src_package ~allowmismatch univ in
  let pool = Depsolver_int.init_pool_univ ~global_constraints univ in
  let globalid = Cudf.universe_size univ in

  (* TODO: this does not yet respect the global_constraints option and will
   * ignore essential *)
  let sdg = Strongdeps.strongdeps_univ univ in

  let tointset l = List.fold_left (fun acc e -> IntSet.add (CudfAdd.vartoint univ e) acc) IntSet.empty l in
  let to_set l = List.fold_right IntSet.add l IntSet.empty in
  (* rule: connect each source package to those source packages which build the
   * binary packages that make their strong dependencies *)
  (* if the dependency closure of any direct dependency has a non-empty
   * intersection with the set of strong dependencies, add a connection. *)
  let partitionofsrc src =
    let strongdeps = Defaultgraphs.PackageGraph.succ_list sdg src in
    let strongdepss = tointset strongdeps in
    let l = List.fold_left (fun acc vpkglist ->
      let disj = CudfAdd.resolve_vpkgs_int univ vpkglist in
      List.fold_left (fun acc2 pid ->
        let dc = Depsolver_int.dependency_closure_cache pool [pid; globalid] in
        let dcs = to_set dc in
        let inters = IntSet.inter strongdepss dcs in
        if not(IntSet.is_empty inters) then (pid, inters)::acc2 else acc2
      ) acc disj
    ) [] src.Cudf.depends in
    IntSet.empty, l
  in

  let instsets = Parmap.parmap ~ncores:num_cores (fun srcpkg ->
    Util.Progress.progress progressbar_build_graph;
    let _, partitions = partitionofsrc srcpkg in
    (srcpkg, partitions)
    ) (Parmap.L sl) in

  Util.Timer.stop timer_build_graph (graph_helper available bin2src univ instsets)
;;

(* Build the dependency source graph bfs from a root node *)
let dist_graph_bfs ?(global_constraints=true) ?(maxdepth=max_int) ?(available=(fun _ -> true)) ?(allowmismatch=false) custom_is_ht univ rootlist =
  Util.Timer.start timer_build_graph_bfs;
  Util.Progress.set_total progressbar_build_graph_bfs (Cudf.universe_size univ);

  let g = G.create () in
  let pool = Depsolver_int.init_pool_univ ~global_constraints univ in

  let bin2src = get_src_package ~allowmismatch univ in

  let queue = Queue.create () in
  let visited = Hashtbl.create (Cudf.universe_size univ) in
  (* add the root source package *)
  List.iter (fun root -> Queue.add (root,0) queue ) rootlist;

  begin try while (Queue.length queue > 0) do
    let (srcpkg,level) = Queue.take queue in
    let src = Unique.create (SrcPkg (CudfAdd.vartoint univ srcpkg)) in
    G.add_vertex g src;
    if not(Hashtbl.mem visited src) && level < maxdepth then begin
      Util.Progress.progress progressbar_build_graph_bfs;
      Hashtbl.add visited src ();
      (* if the package cannot be compiled the List.iter will not happen *)
      let _, partitions = BootstrapCommon.compute_dependency_sets ~global_constraints ~partition:true custom_is_ht pool univ srcpkg in
      List.iter (fun (pid,is) ->
          (* let open Cudf in Cudf_printer.pp_package stdout (CudfAdd.inttovar univ pid); *)
          let instset = Unique.create (InstSet (pid,is)) in
          let label = { depend = BuildDep; annotation = [] } in
          let edge = (src,label,instset) in
          (* src -> instset edge *)
          G.add_edge_e g edge;
          IntSet.iter (fun pid ->
            let binpkg = CudfAdd.inttovar univ pid in
            if not (available binpkg) then begin
              let srcpkgdep = bin2src binpkg in
              let dst = Unique.create (SrcPkg (CudfAdd.vartoint univ srcpkgdep)) in
              (* we add a new src to the queue of sources to compile *)
              Queue.add (srcpkgdep,level+1) queue;
              (* the edge (instset -> src edge) aleady exists. In this case we
               * add the pid of the binary to the list BuildsFrom *)
              try match G.E.label (G.find_edge g instset dst) with
                |{ depend = BuildsFrom s } -> s := IntSet.add pid !s
                |_ -> assert false 
              with Not_found ->  begin
                let label = { depend = BuildsFrom (ref (IntSet.singleton pid)); annotation = [] } in
                let edge = (instset,label,dst) in
                (* instset -> src edge *)
                G.add_edge_e g edge
              end
            end
          ) is
      ) partitions;
    end;
    (* here we check if the graph contains cycles. if this is the case we stop *)
    if false then raise Exit
  done with Exit -> () end ;
  Util.Timer.stop timer_build_graph_bfs g
;;

let from_ic universe ic =
  let getstr n l =
   try match List.assoc n l with
     | GraphmlReader.String s -> s
     | _ -> failwith "expected string"
   with Not_found -> failwith (Printf.sprintf "cannot find key %s in list" n)
  in
  let getint n l =
   try match List.assoc n l with
     | GraphmlReader.Int i -> i
     | _ -> failwith "expected integer"
   with Not_found -> failwith (Printf.sprintf "cannot find key %s in list" n)
  in

  (* this hashtable cuts about 1 second of runtime on a full graph for debian
   * sid with 16k vert and 200k edges *)
  let stringtointmap = Hashtbl.create 97 in

  let stringtoint s =
    match Hashtbl.find_option stringtointmap s with
   | Some i -> i
   | None -> begin
      let name,version = String.split s " " in
      let name = CudfAdd.encode name in
      (* extract 1.0.0 from "(= 1.0.0)" *)
      let version = String.sub version 3 ((String.length version) - 4) in
      let pkgs = List.filter (fun p ->
        CudfAdd.string_of_version p = version
      ) (Cudf.lookup_packages universe name) in
      let pkg = match pkgs with
       | [] -> failwith (Printf.sprintf "cannot find %s (=%s)" name version)
       | [p] -> p
       | _ -> failwith (Printf.sprintf "more than one match for %s (=%s)" name version)
      in
      let i = CudfAdd.vartoint universe pkg in
      Hashtbl.add stringtointmap s i;
      i
    end
  in

  let node l =
    match getstr "kind" l with
     | "SrcPkg" -> begin
       let name = CudfAdd.encode (getstr "name" l) in
       let cudfversion = getint "cudfversion" l in
       let srcpkg = try Cudf.lookup_package universe (name,cudfversion)
       with Not_found -> failwith (Printf.sprintf "cannot find cudf package %s (= %d)" name cudfversion) in
       Unique.create (SrcPkg (CudfAdd.vartoint universe srcpkg))
     end
     | "InstSet" -> begin
       let name = CudfAdd.encode (getstr "name" l) in
       let cudfversion = getint "cudfversion" l in
       let pkg = try Cudf.lookup_package universe (name,cudfversion)
       with Not_found -> failwith (Printf.sprintf "cannot find cudf package %s (= %d)" name cudfversion) in
       let pid = CudfAdd.vartoint universe pkg in
       let binaries = String.nsplit (getstr "binaries" l) "," in
       let is = List.fold_right (fun s -> IntSet.add (stringtoint s)) binaries IntSet.empty in
       Unique.create (InstSet (pid,is))
     end
     | _ -> failwith "invalid node kind"
  in
  let edge l =
    match getstr "kind" l with
     | "builddep" ->
       {depend = BuildDep; annotation = []}
     | "buildsfrom" ->
       let binaries = String.nsplit (getstr "binaries" l) "," in
       let is = List.fold_right (fun s -> IntSet.add (stringtoint s)) binaries IntSet.empty in
       {depend = BuildsFrom (ref is); annotation = []}
     | _ -> failwith "invalid edge kind"
  in

  let module GB = Graph.Builder.I(G) in
  let module GR = GraphmlReader.Parse(GB)(struct let node = node let edge = edge end) in
  (* we want the graph builder to traverse the nodes and edges of the input XML
   * in a deterministic order. This is so that given the same input graph
   * (where "same" means that all node and edge attributes are the same but
   * their order in the XML document may be different) the resulting SrcPkg
   * and InstSet nodes get assigned the same Unique.uid. This in turn is so
   * that it is possible to quickly get a stable order of all vertices
   * between individual runs because many algorithms will sort input
   * vertex lists and sorting by their Unique.uid is many times faster than
   * sorting by package name and version comparison.
   *
   * So instead of letting all algorithms using this graph sort by package name
   * and version, we want to be able to allow quick integer based sorts. But
   * for that the integer assigned to each package name/version must be
   * the same between different runs even if the order in which nodes and
   * edges in the input graph should be different.
   *
   * We rely on the mandatory "id" attribute of nodes and edges to be the same
   * between different runs for the same content. We expect only the order
   * to be different as the order should not matter (both nodes and edges are
   * sets)
   * *)
  let sort l1 l2 =
    let id1 = List.assoc "id" l1 in
    let id2 = List.assoc "id" l2 in
    Pervasives.compare id1 id2
  in
  GR.parse ~nodesort:(Some sort) ~edgesort:(Some sort) ic
;;

module Graphml (U : sig val univ : Cudf.universe end)  = struct
  include G

  let vertex_properties =
      ["name","string",None;
       "version","string",None;
       "binaries","string",None;
       "cudfversion","int",None;
       "architecture","string",None;
       "priority","string",None;
       "type","string",None;
       "essential","string",None;
       "buildessential","string",None;
       "source","string",None;
       "sourcenumber","string",None;
       "kind","string",None;
      ]

  let edge_properties = [
    "kind","string",Some "builddep";
    "vpkglist","string",None;
    "binaries","string",None;
    "annotation","string",None
    ]

  let string_of_vertex_kind = function
    |InstSet _ -> "InstSet"
    |SrcPkg _ -> "SrcPkg"

  let map_vertex vertex =
    let vertex = Unique.value vertex in
    match vertex with
    |SrcPkg id->
      let pkg = CudfAdd.inttovar U.univ id in
      let name = CudfAdd.decode pkg.Cudf.package in
      let version = CudfAdd.string_of_version pkg in

      let kind = ("kind",string_of_vertex_kind vertex) in

      let name = ("name",name) in
      let version = ("version",version) in

      let prop =
        (* don't do anything if key is "version" and translate "cudfversion" to
         * "version" *)
        List.filter_map (fun (key,_,_) ->
          if key = "version" then None else begin
            let k = if key = "cudfversion" then "version" else key in
            try
              let value = Cudf.lookup_package_property pkg k in
              Some(key,value)
            with Not_found -> None
          end
        ) vertex_properties
      in
      let prop = kind :: name :: version :: prop in
      (* only set those attributes which are not empty *)
      List.filter_map (fun (k,v) ->
        if v = "" then None else Some(k,v)
      ) prop
    |InstSet (pid,is) ->
      let binaries =
        let bl =
          List.map (fun pid ->
            let pkg = CudfAdd.inttovar U.univ pid in
            CudfAdd.string_of_package pkg
          ) (IntSet.elements is)
        in
        ("binaries",String.concat "," bl)
      in
      let pkg = CudfAdd.inttovar U.univ pid in
      let name = CudfAdd.decode pkg.Cudf.package in
      let version = CudfAdd.string_of_version pkg in
      let kind = ("kind",string_of_vertex_kind vertex) in
      let name = ("name",name) in
      let version = ("version",version) in
      let prop =
        List.filter_map (fun (key,_,_) ->
          if key = "version" then None else begin
            let k = if key = "cudfversion" then "version" else key in
            try let value = Cudf.lookup_package_property pkg k in
              Some(key,value)
            with Not_found -> None
          end
        ) vertex_properties
      in
      let prop = kind :: name :: version :: binaries :: prop in
      (* only set those attributes which are not empty *)
      List.filter_map (fun (k,v) ->
        if v = "" then None else Some(k,v)
      ) prop

  let map_edge (_,label,_) =
    let annot =
      match label with
        | {annotation} ->
          let annot = List.filter_map (fun a ->
            match a with
              | `StrongDep -> Some("strong")
              | _ -> None
          ) annotation in
          ("annotation", String.concat "," annot)
    in
    match label with
    |{ depend = BuildsFrom { contents = s } } ->
        let witness =
          let sl =
            List.map (fun pid ->
              let pkg = CudfAdd.inttovar U.univ pid in
              CudfAdd.string_of_package pkg
            ) (IntSet.elements s)
          in
          ("binaries", String.concat "," sl)
        in
        let prop = [("kind","buildsfrom");witness;annot] in
        (* only set those attributes which are not empty *)
        List.filter_map (fun (k,v) ->
          if v = "" then None else Some(k,v)
        ) prop
    |{ depend = BuildDep } ->
        let kind = ("kind","builddep") in
        let prop = [kind;annot] in
        (* only set those attributes which are not empty *)
        List.filter_map (fun (k,v) ->
          if v = "" then None else Some(k,v)
        ) prop

  let vertex_uid = G.V.hash
  let edge_uid e = Hashtbl.hash (vertex_uid (G.E.src e),G.E.label e,vertex_uid (G.E.dst e))
end

module Oper = Defaultgraphs.GraphOper(G)
module Comp = Graph.Components.Make(G)
module Cycles = GraphUtils.FindCycles(G)
module Utils = GraphUtils.GraphUtils(G)
module Dfs = Graph.Traverse.Dfs(G)
module T = Graph.Topological.Make(G)
module Printer (U : sig val univ : Cudf.universe end) = Graph.Graphml.Print(G)(Graphml(struct let univ = U.univ end))
