#!/usr/bin/env ocaml
(* {{{ COPYING *)
(*
 * +--------------------------------------------------------------------+
 * | Copyright (C) 2005 2006 Florent Monnier                            |
 * +--------------------------------------------------------------------+
 * | This is a small implementation of the 'tree' command-line utility. |
 * +--------------------------------------------------------------------+
 * |                                                                    |
 * | 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.                       |
 * |                                                                    |
 * | http://www.fsf.org/licensing/licenses/gpl.html                     |
 * |                                                                    |
 * | 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.,                       |
 * | 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA       |
 * |                                                                    |
 * +--------------------------------------------------------------------+
 * | Author: Florent Monnier <fmonnier(€)linux-nantes.fr.eu.org>        |
 * +--------------------------------------------------------------------+
 *
 * }}} *)
#load "unix.cma"

let concat = Filename.concat ;;
(* {{{ usage *)

let usage() =
  print_endline "
  display file times:
    -mt  --modification-time
    -at  --last-access-time
    -ct  --status-change-time

  display with colors:
    -c  --colors
";
  exit(1);
;;
(* }}} *)
(* {{{ round *)

let round nb dec =
  let rec _loop level =
    if level <= 0
    then 1
    else 10 * _loop (pred level)
  in
  let mult = float(_loop dec) in
  let nb = nb *. mult in
  let nb_floor = floor nb
  and nb_ceil  = ceil nb
  in
  if (nb -. nb_floor) < (nb_ceil -. nb)
  then nb_floor /. mult
  else nb_ceil  /. mult
;;
(* }}} *)
(* {{{ string_of_round *)

let string_of_round_nopad nb dec =
  let rounded = round nb dec in
  if rounded = (floor nb)
  then string_of_int (int_of_float nb)
  else string_of_float rounded
;;

let string_of_round nb dec =
  let rounded = round nb dec in
  if rounded = (floor nb)
  then Printf.sprintf "%4d" (int_of_float nb)
  else Printf.sprintf "%4.f" rounded
;;
(* }}} *)
(* {{{ human_size *)

let units =
  let u = Array.make 4 " " in
  u.(1) <- "K";
  u.(2) <- "M";
  u.(3) <- "G";
  u
;;

let human_size ~size =
  let rec aux size level =
    if size < 1024.0 || level = 3
    then (string_of_round size 1), units.(level)
    else aux (size /. 1024.0) (succ level)
  in
  let size = float_of_int size in
  let (size, _unit) = aux size 0 in
  if _unit = " "
  then Printf.sprintf "[%s%s]" _unit size
  else Printf.sprintf "[%s%s]" size _unit
;;

let units =
  let u = Array.make 4 " " in
  u.(1) <- "K";
  u.(2) <- "M";
  u.(3) <- "G";
  u
;;

let human_size ~size =
  let rec aux size level =
    if size < 1024.0 || level = 3
    then Printf.sprintf "%6.1f%s" size units.(level)
    else aux (size /. 1024.0) (succ level) in
  let size = float_of_int size in
  aux size 0
;;
(* }}} *)
(* {{{ padding *)

let padding ~last ~depth =
  let rec aux  depth last pad_acc =
    if depth <= 0
    then pad_acc
    else begin
      let branch =
        match last with
        | false :: _ -> "|   "
        | true  :: _ -> "    "
        | [] -> "xxxx"  (* this case should never occur *)
      in
      aux (pred depth) (List.tl last) (branch ^ pad_acc)
    end
  in
  let last =
    match last with
    | head::tail -> tail  (* the head is printed from branch_mark *)
    | _ -> []
  in
  aux depth last ""
;;


let right_pad filenames =
  let rec pass_1 width rev = function [] -> (rev, width)
    | (filename,s)::t ->
        let len = String.length filename in
        pass_1 (max width len) ((filename,len,s)::rev) t
  in
  let rev, width = pass_1 0 [] filenames in

  let rec pass_2 padded = function [] -> padded
    | (filename,len,s)::t ->
        let n = width - len in
        let pad = String.make n ' ' in
        let str = (filename ^ pad) in
        pass_2 ((str,s)::padded) t
  in
  pass_2 [] rev
;;

(* }}} *)
(* {{{ branch_mark *)

let branch_mark ~last =
  match last with
  | false :: _ -> "|--"
  | true  :: _ -> "`--"
  | [] -> "" (* only from boot_dump for base_dir *)
;;

let branch_dir_mark ~last =
  match last with
  | false :: _ -> "+-- "
  | true  :: _ -> "`-- "
  | [] -> "" (* only from boot_dump for base_dir *)
;;
(* }}} *)
(* {{{ color *)
let g = char_of_int 27 ;;

let color color_name ?(label="") str () =
  let col_code = match color_name with
  (*
  | `blue      -> "01;34"
  | `yellow    -> "01;33"
  *)
  | `dark_red  -> "02;31"
  | `purple    -> "03;35"
  | `dark_grey -> "01;30"
  | `test      -> "01;32"

  |     `red -> "31;49"
  |   `green -> "32;49"
  |  `yellow -> "33;49"
  |    `blue -> "34;49"
  | `magenta -> "35;49"
  |    `cyan -> "36;49"
  |   `white -> "37;49"
  | `default -> "39;49"
  in
  Printf.sprintf "%c[%sm%s%s%c[00m" g col_code  label str  g
;;
(* }}} *)
(* {{{ human_perms *)

let human_perms ~perms =
  let string_of_octal = function
    | '1' -> "--x"
    | '2' -> "-w-"
    | '3' -> "-wx"
    | '4' -> "r--"
    | '5' -> "r-x"
    | '6' -> "rw-"
    | '7' -> "rwx"
    |  _  -> "---"
  in
  let octal_str =
    Printf.sprintf "%04o" perms
  in
  let u = string_of_octal  octal_str.[1]
  and g = string_of_octal  octal_str.[2]
  and o = string_of_octal  octal_str.[3]
  in
  (u ^ g ^ o)
;;
(* }}} *)
(* {{{ dump_total_size *)

let dump_total_size ~size ~depth ~last ~colors =
  if size = 0 then () else begin
    let pad = padding ~last:(true::last) ~depth in
    let h_size = human_size ~size in

    if colors then begin
      Printf.printf "%s" (color `yellow (pad ^ "o--->  ") ());
      Printf.printf "%s" (color `dark_red h_size ());
    end else begin
      Printf.printf "%s" (pad ^ "o--->  ");
      Printf.printf "%s" h_size;
    end;
    Printf.printf "\n";
  end
;;
(* }}} *)
(* {{{ human_time *)

let human_time time =
  (* Unix.gmtime  UTC (Coordinated Universal Time), aka GMT *)
  let t = Unix.localtime time in (* local time zone *)
  Printf.sprintf "%d-%02d-%02d/%02d:%02d"
     (t.Unix.tm_year + 1900)
      t.Unix.tm_mon
      t.Unix.tm_mday
      t.Unix.tm_hour
      t.Unix.tm_min
;;
(* }}} *)
(* {{{ which times *)

let (mt, at, ct) =
  let argc = Array.length Sys.argv in
  let (mt, at, ct) =
    (ref false, ref false, ref false)
  in
  for i = 1 to pred argc do
    if Sys.argv.(i) = "-h"  || Sys.argv.(i) = "--help"  then usage();
    if Sys.argv.(i) = "-mt" || Sys.argv.(i) = "--modification-time"  then mt := true;
    if Sys.argv.(i) = "-at" || Sys.argv.(i) = "--last-access-time"   then at := true;
    if Sys.argv.(i) = "-ct" || Sys.argv.(i) = "--status-change-time" then ct := true;
  done;
  (!mt, !at, !ct)
;;

(* }}} *)
(* {{{ dump_file *)

let dump_file ~name ~stats ~depth ~last ~colors =
  let file_name  = name in

  let perms = stats.Unix.st_perm
  and size  = stats.Unix.st_size
  in
  let atime = stats.Unix.st_atime   (* Last access time *)
  and mtime = stats.Unix.st_mtime   (* Last modification time *)
  and ctime = stats.Unix.st_ctime   (* Last status change time *)
  in
  let h_size = human_size ~size in
  let pad = padding ~last ~depth in

  let show_times = true in

  if colors then begin
    Printf.printf "%s" (color `yellow (pad ^ (branch_mark ~last)) ());
    Printf.printf " %s" (color `purple ~label:"-" (human_perms ~perms) ());
  (*Printf.printf "%s " (color `purple ~label:"perms:" (Printf.sprintf "%03o" perms) ()); (* Octal *) *)
    Printf.printf " %s"  (color `dark_red h_size ());
    Printf.printf " %s" (Filename.basename file_name);
    if show_times then begin
      if mt then Printf.printf "  %s" (color `dark_grey ~label:"mt:" (human_time mtime) ());
      if at then Printf.printf "  %s" (color `dark_grey ~label:"at:" (human_time atime) ());
      if ct then Printf.printf "  %s" (color `dark_grey ~label:"ct:" (human_time ctime) ());
    end;
  end else begin
    Printf.printf "%s" (pad ^ (branch_mark ~last));
    Printf.printf " -%s" (human_perms ~perms);
  (*Printf.printf "perms:%s " (Printf.sprintf "%03o" perms); (* Octal *) *)
    Printf.printf " %s"  (h_size);
    Printf.printf " %s" (Filename.basename file_name);
    if show_times then begin
      if mt then Printf.printf "  %s" ("mt:" ^ (human_time mtime));
      if at then Printf.printf "  %s" ("at:" ^ (human_time atime));
      if ct then Printf.printf "  %s" ("ct:" ^ (human_time ctime));
    end;
  end;
  Printf.printf "\n";
  (size)
;;
(* Unix.stats.Unix.st_mtime *)
(* }}} *)
(* {{{ dump_file_list *)

let dump_file_list  ~files ~dirs ~parent_dir ~depth ~last ~colors =

  let rec file_loop ~files ~dirs size_acc =
    match files, dirs with
    | [], [] | [], _ ->   (* this case matches a directory containing sub-directories *)
        dump_total_size ~size:(size_acc) ~depth ~last ~colors;
        (size_acc)

    | (file_name, stats)::[], [] ->  (* this case matches a file leaf *)
        let file_size = dump_file
          ~name:(concat parent_dir file_name) ~stats ~depth ~last:(true::last) ~colors
        in
        let total_size = file_size + size_acc in
        dump_total_size ~size:(total_size) ~depth ~last ~colors;
        (total_size)

    | (file_name, stats)::tail_files, dirs ->  (* iteration loop *)
        let file_size = dump_file
          ~name:(concat parent_dir file_name) ~stats ~depth ~last:(false::last) ~colors
        in
        file_loop ~files:tail_files ~dirs (file_size + size_acc)
  in
  file_loop ~files ~dirs 0;
;;

(* }}} *)
(* {{{ dump_dir[_content] *)

let rec dump_dir ~name ~stats ~depth ~last ~colors =
  let parent_dir  = name
  and parent_dir_stats = stats in

  let pad = padding ~last ~depth:(pred depth) in

  if colors then
    Printf.printf "%s%s\n"  (* with colors *)
        (color `yellow (pad ^ (branch_dir_mark ~last)) ())
        (color `blue (parent_dir ^ "/") ())
  else
    Printf.printf "%s%s%s/\n" pad (branch_dir_mark ~last) parent_dir;  (* without colors *)

  let contents = Sys.readdir parent_dir in
  let contents = Array.to_list contents in
  let rec sort c ~reg_acc ~dir_acc =
    match c with
    | [] -> (reg_acc, dir_acc)
    | name :: tl ->
        let stats = Unix.lstat (concat parent_dir name) in
        match stats.Unix.st_kind with
        (* Accumulate directories and regular files in 2 different lists: dir_acc & reg_acc *)
        | Unix.S_REG -> sort tl ~dir_acc ~reg_acc:((name,stats) :: reg_acc)
        | Unix.S_DIR -> sort tl ~reg_acc ~dir_acc:((name,stats) :: dir_acc)
        | Unix.S_LNK | Unix.S_CHR | Unix.S_BLK | Unix.S_FIFO | Unix.S_SOCK ->
            sort tl ~dir_acc ~reg_acc (* XXX *)
  in
  let (files, dirs) = sort contents ~reg_acc:[] ~dir_acc:[] in
  let files = right_pad files in

  (* For a better readablility, the files are printed before directories contents. *)
  let files_size = dump_file_list ~files ~dirs ~parent_dir ~depth ~last ~colors in

  (* {{{ print directories *)
  let rec dir_loop dirs size_acc =
    match dirs with
    | [] -> size_acc
    | (dir_name, stats) :: [] ->   (* this case matches the last directory *)
        let content_size = dump_dir
          ~name:(concat parent_dir dir_name) ~stats ~depth:(succ depth) ~last:(true::last) ~colors;
        in
        (content_size + size_acc)

    | (dir_name, stats) :: tail_dirs ->
        let content_size = dump_dir
          ~name:(concat parent_dir dir_name) ~stats ~depth:(succ depth) ~last:(false::last) ~colors;
        in
        dir_loop tail_dirs (content_size + size_acc)
  in
  let rec_content_size = dir_loop dirs files_size in
  (rec_content_size)
  (* }}} *)
;;

(* }}} *)
(* {{{ main *)

let boot_dump  base_dir  options =
  let stats = Unix.lstat base_dir in
  let colors =
    List.mem "--colors" options || List.mem "-c" options
  in
  let all_size = dump_dir ~name:base_dir ~stats ~depth:0 ~last:[] ~colors in
  if colors then
    Printf.printf "\n%s\n" (color `green ~label:"Total size:" (human_size all_size) ())
  else
    Printf.printf "\nTotal size:%s\n" (human_size all_size);
;;

let () =
  let argc = Array.length Sys.argv in

  let this_dir d yet options =
    if Sys.file_exists d
    then (boot_dump d options; print_newline(); succ yet, options)
    else ((*Printf.printf "Warning: '%s' does not exist\n" d;*) yet, d::options)
  in
  let rec parse_arg argi yet options =
    if argi < argc then begin
      let (yet, options) =
        this_dir Sys.argv.(argi) yet options
      in
      parse_arg (succ argi) yet options
    end else if yet = 0 then
      ignore(this_dir "." 0 options) (* Sys.getcwd () *)
  in
  parse_arg 1 0 [];
;;
(* }}} *)

(* vim:cindent sw=2 sts=2 ts=2 et fdm=marker
 *)