(* rpmdepsize - visualize the size of RPM dependencies
 * (C) Copyright 2009 Red Hat Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 * Written by Richard W.M. Jones <rjones@redhat.com>
 *)

open Sexplib
TYPE_CONV_PATH "."

open ExtList
open Unix
open Printf

let debug = true

(* This corresponds to the sexpr that we write out from the
 * Python code.  OCaml will type-check it.
 *)
type root_packages = string * packages
and packages = pkg list
and pkg = {
  nevra : string;			(* name-[epoch:]version-release.arch *)
  name : string;
  epoch : int;
  version : string;
  release : string;
  arch : string;
  size : int64;				(* installed size, excl. dirs *)
  deps : string list;
}
 with sexp

(* Full dependency representation.  This is actually a graph because
 * it contains dependency loops.  'deps list' is a ref because we
 * update it as we are building it.
 *)
type deps = Deps of pkg * deps list ref

(* Final tree representation, loops removed, and everything we want to
 * display stored in the nodes.
 *)
type tree = Tree of pkg * int64 * int64 * GDraw.color * tree list

module StringMap = Map.Make (String)
let (+^) = Int64.add
let sum = List.fold_left (+^) 0L
let spaces n = String.make n ' '

(* Python has privileged access to the yum repodata, so we have to use
 * this Python snippet to pull the data that we need out.  This is the
 * part of the program that takes ages to run, because Python is as
 * slow as a fat snake that's just eaten a huge lunch.  We can't help that.
 *)
let repoquery_py = "
import yum
import yum.misc
import sys

yb = yum.YumBase ()

basepkg = yb.pkgSack.returnPackages (patterns=[sys.argv[1]])[0]
deps = dict ({basepkg:False})

# Recursively find all the dependencies.
stable = False
while not stable:
    stable = True
    for pkg in deps.keys():
        if deps[pkg] == False:
            deps[pkg] = []
            stable = False
            for r in pkg.requires:
                ps = yb.whatProvides (r[0], r[1], r[2])
                best = yb._bestPackageFromList (ps.returnPackages ())
                if best.name != pkg.name:
                    deps[pkg].append (best)
                    if not deps.has_key (best):
                        deps[best] = False
            deps[pkg] = yum.misc.unique (deps[pkg])

# Get the data out of python as fast as possible so we can
# use a serious language for analysis of the tree.
print \"(%s (\" % basepkg
for pkg in deps.keys():
    print \"((nevra %s) (name %s) (epoch %s) (version %s) (release %s) (arch %s) (size %s)\" % (pkg, pkg.name, pkg.epoch, pkg.version, pkg.release, pkg.arch, pkg.installedsize)
    print \"(deps (\"
    for p in deps[pkg]:
        print \"%s \" % p,
    print \")))\"
sys.stdout.write (\"))\")  # suppress trailing newline"

let () =
  printf "getting repository information (this can take a few seconds ...)\n%!";

  (* Run the Python program and read in the generated sexpr. *)
  let cmd =
    sprintf "python -c %s %s"
      (Filename.quote repoquery_py) (Filename.quote Sys.argv.(1)) in
  let chan = open_process_in cmd in
  ignore (input_line chan); (* drop "Loaded plugins" *)
  let root, pkgs =
    root_packages_of_sexp (Sexp.of_string (Std.input_all chan)) in
  (match close_process_in chan with
   | WEXITED 0 -> ()
   | WEXITED i -> failwith (sprintf "command exited with status %d" i)
   | WSIGNALED i | WSTOPPED i ->
       failwith (sprintf "command stopped with signal %d" i)
  );

  (* Create the dependency graph, probably contains loops so beware. *)
  let deps = List.map (fun pkg -> Deps (pkg, ref [])) pkgs in
  let depsmap =
    List.fold_left (
      fun map (Deps (pkg, _) as deps) ->
	StringMap.add pkg.nevra deps map
    ) StringMap.empty deps in
  List.iter (
    fun (Deps (pkg, deps)) ->
      let deps' = List.map (fun n -> StringMap.find n depsmap) pkg.deps in
      deps := List.append !deps deps'
  ) deps;

  (* For each package, calculate the total installed size of the package,
   * which includes all subpackages pulled in.  So it's what would be
   * installed if you did 'yum install foo'.
   *)
  let total pkg =
    let seen = ref StringMap.empty in
    let rec _total = function
      | Deps (pkg, _) when StringMap.mem pkg.nevra !seen -> 0L
      | Deps (pkg, { contents = children }) ->
	  seen := StringMap.add pkg.nevra true !seen;
	  pkg.size +^ sum (List.map _total children)
    in
    _total (StringMap.find pkg.nevra depsmap)
  in
  let totalsmap =
    List.fold_left (
      fun map pkg -> StringMap.add pkg.nevra (total pkg) map
    ) StringMap.empty pkgs in

  (* Create the final display tree.  Each node is sorted so that
   * children with the largest contribution come first (on the left).
   * We remove packages which are already installed by earlier
   * (leftward) packages.  At each node we also store total size and
   * size of the additional packages.
   *)
  let tree =
    let seen = ref StringMap.empty in
    let rec build_tree = function
      | Deps (pkg, _) when StringMap.mem pkg.nevra !seen -> None
      | Deps (pkg, { contents = children }) ->
	  (* Sort children by reverse total size. *)
	  let cmp (Deps (p1, _)) (Deps (p2, _)) =
	    let t1 = StringMap.find p1.nevra totalsmap in
	    let t2 = StringMap.find p2.nevra totalsmap in
	    compare t2 t1
	  in
	  let children = List.sort ~cmp children in
	  seen := StringMap.add pkg.nevra true !seen;
	  let children = List.filter_map build_tree children in
	  let total = StringMap.find pkg.nevra totalsmap in
	  let increm =
	    let rec sum_child_sizes = function
	      | Tree (pkg, _, _, _, children) ->
		  List.fold_left (
		    fun size child -> size +^ sum_child_sizes child
		  ) pkg.size children
	    in
	    sum_child_sizes (Tree (pkg, 0L, 0L, `WHITE, children)) in
	  Some (Tree (pkg, total, increm, `WHITE, children))
    in
    Option.get (build_tree (StringMap.find root depsmap)) in

  if debug then (
    let rec display ?(indent=0) = function
      | Tree (pkg, total, increm, _, children) ->
	  printf "%s%s %Ld/%Ld/%Ld\n%!"
	    (spaces indent) pkg.nevra pkg.size increm total;
	  List.iter (display ~indent:(indent+2)) children
    in
    display tree
  );

  (* Max depth of the tree. *)
  let depth =
    let rec depth = function
      | Tree (pkg, _, _, _, children) ->
	  List.fold_left (fun d c -> max d (1 + depth c)) 1 children
    in
    depth tree in

  (* Allocate a colour to each node in the tree based on its parent.  The
   * single top node is always light grey.  The second level nodes are
   * primary colours.
   *)
  let tree =
    let Tree (pkg, total, increm, _, level2) = tree in
    let level2 =
      let pcols = [
	`RGB (55000, 0, 0);
	`RGB (0, 55000, 0);
	`RGB (0, 0, 55000);
	`RGB (55000, 55000, 0);
	`RGB (0, 55000, 55000);
      ] in
      let rec colour_level2 cols = function
	| [] -> []
	| Tree (pkg, total, increm, _, level3) :: level2 ->
	    let col, cols = match cols with
	      | [] -> List.hd pcols, List.tl pcols
	      | col :: cols -> col, cols in
	    let level3 = colour_level3 col (List.length level3) 0 level3 in
	    Tree (pkg, total, increm, col, level3)
	    :: colour_level2 cols level2
      and colour_level3 col n i = function
	| [] -> []
	| Tree (pkg, total, increm, _, leveln) :: level3 ->
	    let col = scale_colour col n i in
	    let leveln = colour_level3 col (List.length leveln) 0 leveln in
	    Tree (pkg, total, increm, col, leveln)
	    :: colour_level3 col n (i+1) level3
      and scale_colour col n i =
	let r, g, b = match col with
	  | `RGB (r, g, b) -> float r, float g, float b
	  | _ -> assert false in
	let i = float i and n = float n in
	let scale = 0.8 +. i/.(5.*.n) in
	let r = r *. scale in
	let g = g *. scale in
	let b = b *. scale in
	`RGB (int_of_float r, int_of_float g, int_of_float b)
      in
      colour_level2 pcols level2 in
    Tree (pkg, total, increm, `RGB (55000, 55000, 55000), level2) in

  (* Open the window. *)
  let title = root ^ " - Fedora RPM dependency size viewer" in
  let window =
    GWindow.window ~width:800 ~height:600 ~title ~allow_shrink:true () in

  ignore (window#connect#destroy ~callback:GMain.quit);

  let da = GMisc.drawing_area ~packing:window#add () in
  da#misc#realize ();
  let draw = new GDraw.drawable da#misc#window in

  (* Pango contexts used to draw large and small text. *)
  let pango_large_context = da#misc#create_pango_context in
  pango_large_context#set_font_description (Pango.Font.from_string "Sans 12");
  let pango_small_context = da#misc#create_pango_context in
  pango_small_context#set_font_description (Pango.Font.from_string "Sans 8");

  let repaint _ =
    (* Get the canvas size and fill the background with white. *)
    let width, height = draw#size in
    draw#set_background `WHITE;
    draw#set_foreground `WHITE;
    draw#rectangle ~x:0 ~y:0 ~width ~height ~filled:true ();

    (* Calculate the scales so we can fit everything into the window. *)
    let Tree (_, top_total, top_increm, _, _) = tree in
    assert (top_total = top_increm);
    let rowheight = float height /. float depth in
    let scale = float width /. Int64.to_float top_increm in

    (* Now draw the tree. *)
    let rec draw_tree y x = function
      | Tree (pkg, total, increm, colour, children) ->
	  (* Draw pkg at (x, y). *)
	  let width = scale *. Int64.to_float increm in
	  let pkgsizewidth = scale *. Int64.to_float pkg.size in
	  draw_pkg x y width pkgsizewidth rowheight colour pkg total increm;

	  (* Draw the children of pkg at (i, y + rowheight), where
	   * i starts as x and increments for each child.
	   *)
	  let y = y +. rowheight in
	  let rec loop x = function
	    | [] -> ()
	    | child :: children ->
		draw_tree y x child;
		let Tree (_, _, increm, _, _) = child in
		let childwidth = scale *. Int64.to_float increm in
		loop (x +. childwidth) children
	  in
	  loop x children

    (* Draw a single package. *)
    and draw_pkg x y width pkgsizewidth height colour pkg total increm =
      let x = int_of_float x in
      let y = int_of_float y in
      let width = int_of_float width in
      let pkgsizewidth = int_of_float pkgsizewidth in
      let height = int_of_float height in

      if width > 8 then (
	draw_pkg_outline x y width pkgsizewidth height colour;
	draw_pkg_label x y width height colour pkg total increm
      )
      else if width >= 4 then
	draw_pkg_narrow x y width height colour
      (* else nothing *)

    and draw_pkg_outline x y width pkgsizewidth height colour =
      draw#set_foreground colour;
      draw#rectangle ~x:(x+2) ~y:(y+2)
	~width:(width-4) ~height:(height-4)
	~filled:true ();
      if pkgsizewidth > 2 then (
	draw#set_foreground (darken colour);
	draw#rectangle ~x:(x+2) ~y:(y+2)
	  ~width:(pkgsizewidth-2) ~height:(height-4)
	  ~filled:true ();
	draw#set_foreground (choose_contrasting_colour colour);
	draw#set_line_attributes ~style:`ON_OFF_DASH ();
	draw#line (x+pkgsizewidth) (y+2) (x+pkgsizewidth) (y+height-2);
	draw#set_line_attributes ~style:`SOLID ()
      );
      draw#set_foreground (`BLACK);
      draw#rectangle ~x:(x+2) ~y:(y+2)
	~width:(width-4) ~height:(height-4)
	~filled:false ()

    and draw_pkg_label x y width height colour pkg total increm =
      (* How to write text in a drawing area, in case it's not
       * obvious, which it certainly is not:
       * http://www.math.nagoya-u.ac.jp/~garrigue/soft/olabl/lablgtk-list/120.txt
       *)
      let txt1 = lazy (
	let txt = pango_large_context#create_layout in
	Pango.Layout.set_text txt (
	  sprintf "%s
Package: %.1f%% %s (%Ld bytes)
Incremental: %.1f%% %s (%Ld bytes)
Total: %.1f%% %s (%Ld bytes)" pkg.nevra
(display_percent pkg.size) (display_size pkg.size) pkg.size
(display_percent increm) (display_size increm) increm
(display_percent total) (display_size total) total
	);
	txt
      )
      and txt2 = lazy (
	let txt = pango_small_context#create_layout in
	Pango.Layout.set_text txt (
	  sprintf "%s
Package: %.1f%% %s (%Ld bytes)
Incremental: %.1f%% %s (%Ld bytes)
Total: %.1f%% %s (%Ld bytes)" pkg.nevra
(display_percent pkg.size) (display_size pkg.size) pkg.size
(display_percent increm) (display_size increm) increm
(display_percent total) (display_size total) total
	);
	txt
      )
      and txt3 = lazy (
	let txt = pango_small_context#create_layout in
	Pango.Layout.set_text txt (
	  sprintf "%s
Pkg: %.1f%% %s (%Ld bytes)
Incr: %.1f%% %s (%Ld bytes)
Tot: %.1f%% %s (%Ld bytes)" pkg.name
(display_percent pkg.size) (display_size pkg.size) pkg.size
(display_percent increm) (display_size increm) increm
(display_percent total) (display_size total) total
	);
	txt
      )
      and txt4 = lazy (
	let txt = pango_small_context#create_layout in
	Pango.Layout.set_text txt (
	  sprintf "%s
Pkg: %.1f%% %s
Incr: %.1f%% %s
Tot: %.1f%% %s" pkg.name
(display_percent pkg.size) (display_size pkg.size)
(display_percent increm) (display_size increm)
(display_percent total) (display_size total)
	);
	txt
      )
      and txt5 = lazy (
	let txt = pango_small_context#create_layout in
	Pango.Layout.set_text txt (
	  sprintf "%s\nPkg: %.1f%%\nIncr: %.1f%%\nTot: %.1f%%"
	    pkg.name
	    (display_percent pkg.size)
	    (display_percent increm)
	    (display_percent total)
	);
	txt
      )
      and txt6 = lazy (
	let txt = pango_small_context#create_layout in
	Pango.Layout.set_text txt (
	  sprintf "%s Pkg: %.1f%% %s Incr: %.1f%% %s Tot: %.1f%% %s" pkg.name
	    (display_percent pkg.size) (display_size pkg.size)
	    (display_percent increm) (display_size increm)
	    (display_percent total) (display_size total)
	);
	txt
      )
      and txt7 = lazy (
	let txt = pango_small_context#create_layout in
	Pango.Layout.set_text txt (
	  sprintf "%s %.1f%% %.1f%% %.1f%%" pkg.name
	    (display_percent pkg.size)
	    (display_percent increm)
	    (display_percent total)
	);
	txt
      )
      and txt8 = lazy (
	let txt = pango_small_context#create_layout in
	Pango.Layout.set_text txt (
	  sprintf "%s" pkg.name
	);
	txt
      ) in
      let txts = [ txt1; txt2; txt3; txt4; txt5; txt6; txt7; txt8 ] in

      let fore = choose_contrasting_colour colour in

      let rec loop = function
	| [] -> ()
	| txt :: txts ->
	    let txt = Lazy.force txt in
	    let { Pango.width = txtwidth;
		  Pango.height = txtheight } =
	      Pango.Layout.get_pixel_extent txt in
	    (* Now with added fudge-factor. *)
	    if width >= txtwidth + 8 && height >= txtheight + 8 then
	      draw#put_layout ~x:(x+4) ~y:(y+4) ~fore txt
	    else loop txts
      in
      loop txts

    and draw_pkg_narrow x y width height colour =
      draw#set_foreground colour;
      draw#rectangle ~x:(x+2) ~y:(y+2)
	~width:(width-4) ~height:(height-4) ~filled:true ()

    and choose_contrasting_colour = function
      | `RGB (r, g, b) ->
	  if r + g + b > 98304 then `BLACK else `WHITE
      | _ -> `WHITE

    and darken = function
      | `RGB (r, g, b) ->
	  `RGB (r * 9 / 10, g * 9 / 10, b * 9 / 10)
      | _ -> `WHITE

    and display_percent bytes =
      100. *. Int64.to_float bytes /. Int64.to_float top_total

    and display_size bytes =
      if bytes > 104_857L then
	sprintf "%.1f MB" (Int64.to_float bytes /. 1_048_576.)
      else if bytes > 102L then
	sprintf "%.1f KB" (Int64.to_float bytes /. 1_024.)
      else
	sprintf "%Ld" bytes
    in
    draw_tree 0. 0. tree;

    (* Return false because this is a Gtk event handler. *)
    false
  in
  ignore (da#event#connect#expose ~callback:repaint);

  window#show ();
  GMain.main ()