(* Basic Merjis functions.
 * Copyright (C) 2004 Merjis Ltd.
 * Written By Richard W.M. Jones (rich@merjis.com)
 * $Id: merjisforwiki.ml,v 1.1 2004/09/07 10:14:09 rich Exp $
 *)

(* This is a copy of an internal library which we use at Merjis. *)

open Printf
open ExtString
open ExtList

(*----- Basic stuff. -----*)

let identity x = x

let unique =
  let n = ref 0 in
  fun () -> incr n; !n

let rec range a b =
  if a <= b then
    a :: range (a+1) b
  else
    []

(*----- String functions. -----*)

let string_contains substr str =
  try String.find str substr; true
  with String.Invalid_string -> false

let string_of_char = String.make 1

let truncate n str =
  if String.length str < n then str else String.sub str 0 (n-1)

(* These versions only work in the C locale for 7-bit characters. *)
let isspace c =
  c = ' '
  (* || c = '\f' *) || c = '\n' || c = '\r' || c = '\t' (* || c = '\v' *)

let isalpha c =
  c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z'

let isdigit c =
  c >= '0' && c <= '9'

let isalnum c =
  c >= '0' && c <= '9' || c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z'

let islower c =
  c >= 'a' && c <= 'z'

let isupper c =
  c >= 'A' && c <= 'Z'

let isxdigit c =
  c >= '0' && c <= '9' || c >= 'a' && c <= 'f' || c >= 'A' && c <= 'F'

let triml ?(test = isspace) str =
  let i = ref 0 in
  let n = ref (String.length str) in
  while !n > 0 && test str.[!i]; do
    decr n;
    incr i
  done;
  if !i = 0 then str
  else String.sub str !i !n

let trimr ?(test = isspace) str =
  let n = ref (String.length str) in
  while !n > 0 && test str.[!n-1]; do
    decr n
  done;
  if !n = String.length str then str
  else String.sub str 0 !n

let trim ?(test = isspace) str =
  trimr (triml str)

let string_for_all f str =
  let len = String.length str in
  let rec loop i =
    if i = len then true
    else (
      let c = str.[i] in
      if not (f c) then false
      else loop (i+1)
    )
  in
  loop 0

let string_exists f str =
  let len = String.length str in
  let rec loop i =
    if i = len then false
    else (
      let c = str.[i] in
      if f c then true
      else loop (i+1)
    )
  in
  loop 0

let string_is_whitespace = string_for_all isspace

(*----- List functions. -----*)

let first n xs =
  fst (List.split_nth n xs)

let rec uniq ?(cmp = Pervasives.compare) = function
    [] -> []
  | [x] -> [x]
  | x :: y :: xs when compare x y = 0 ->
      uniq (x :: xs)
  | x :: y :: xs ->
      x :: uniq (y :: xs)

let sort_uniq ?cmp xs =
  let xs = List.sort ?cmp xs in
  let xs = uniq ?cmp xs in
  xs

let frequency ?(cmp = Pervasives.compare) xs =
  let xs = List.sort ~cmp xs in
  let rec loop = function
      [] -> []
    | [x] -> [1, x]
    | x :: y :: xs when cmp x y = 0 ->
        let rest = loop (y :: xs) in
        let (count, _), rest = List.hd rest, List.tl rest in
        (count+1, y) :: rest
    | x :: xs ->
        (1, x) :: loop xs
  in
  let xs = loop xs in
  List.rev (List.sort xs)

(* This version by Isaac Trotts. *)
let group_by ?(cmp = Pervasives.compare) ls =
  let ls' =
    List.fold_left
      (fun acc (day1, x1) ->
         match acc with
             [] -> [day1, [x1]]
           | (day2, ls2) :: acctl ->
               if cmp day1 day2 = 0
               then (day1, x1 :: ls2) :: acctl
               else (day1, [x1]) :: acc)
      []
      ls
  in
  let ls' = List.rev ls' in
  List.map (fun (x, xs) -> x, List.rev xs) ls'

(*----- File functions. -----*)

let (//) = Filename.concat

let rec input_all_lines chan =
  try
    let line = input_line chan in
    line :: input_all_lines chan
  with
      End_of_file -> []

let input_all chan =
  let buf = Buffer.create 16384 in
  let tmpsize = 16384 in
  let tmp = String.create tmpsize in
  let n = ref 0 in
  while n := input chan tmp 0 tmpsize; !n > 0 do
    Buffer.add_substring buf tmp 0 !n;
  done;
  Buffer.contents buf

let input_file filename =
  let chan = open_in_bin filename in
  let data = input_all chan in
  close_in chan;
  data

let output_file filename data =
  let chan = open_out_bin filename in
  output_string chan data;
  close_out chan

let output_tempfile data =
  let filename, chan = Filename.open_temp_file "tmp" ".tmp" in
  output_string chan data;
  close_out chan;
  filename

(*----- Command functions. -----*)

let cmd cmd =
  let code = Sys.command cmd in
  if code <> 0 then failwith (cmd ^ ": error code " ^ string_of_int code)

let copy infile outfile =
  cmd (sprintf "cp %s %s" infile outfile)

let pget cmd =
  let chan = Unix.open_process_in cmd in
  let lines = input_all_lines chan in
  let stat = Unix.close_process_in chan in
  (match stat with
       Unix.WEXITED 0 -> ()
     | Unix.WEXITED i ->
	 failwith ("command failed with code " ^ string_of_int i)
     | Unix.WSIGNALED i ->
	 failwith ("command killed by signal " ^ string_of_int i)
     | Unix.WSTOPPED i ->
	 failwith ("command stopped by signal " ^ string_of_int i));
  lines

let unlink file =
  try Unix.unlink file with Unix.Unix_error _ -> ()

(*----- Meta-functions. -----*)
let notf f =
  fun v -> not (f v)

(*----- Time and date. -----*)
let short_weekday = function
  | 0 -> "Sun" | 1 -> "Mon" | 2 -> "Tue" | 3 -> "Wed"
  | 4 -> "Thu" | 5 -> "Fri" | 6 -> "Sat" | 7 -> "Sun"
  | _ -> invalid_arg "short_weekday"

let short_month = function
  | 1 -> "Jan" | 2 -> "Feb" | 3 -> "Mar" | 4 -> "Apr"
  | 5 -> "May" | 6 -> "Jun" | 7 -> "Jul" | 8 -> "Aug"
  | 9 -> "Sep" | 10 -> "Oct" | 11 -> "Nov" | 12 -> "Dec"
  | _ -> invalid_arg "short_month"

(*----- Images. -----*)

(* Find the format of an image.  Uses the external 'identify' program,
 * part of ImageMagick.  Returns (mime_type, width, height).  Throws
 * Invalid_argument "image_identify" if the data is not an image.
 *)
let image_identify_re = Pcre.regexp "^\\S+ ([A-Z]+) (\\d+)x(\\d+)"

let image_identify data =
  let filename = output_tempfile data in
  let in_chan = Unix.open_process_in ("identify " ^ filename) in
  let line = input_line in_chan in
  let status = Unix.close_process_in in_chan in
  unlink filename;
  (match status with
       Unix.WEXITED 0 ->		(* Identify was OK with it ... *)
	 ()
     | Unix.WEXITED _ ->		(* Couldn't identify the file type. *)
	 raise (Invalid_argument "image_identify")
     | Unix.WSIGNALED n ->
	 failwith ("image_identify: 'identify' killed by signal " ^
		   string_of_int n)
     | Unix.WSTOPPED n ->
	 failwith ("image_identify: 'identify' stopped by signal " ^
		   string_of_int n));
  try
    let subs = Pcre.exec ~rex:image_identify_re line in
    let type_string = Pcre.get_substring subs 1 in
    let width = int_of_string (Pcre.get_substring subs 2) in
    let height = int_of_string (Pcre.get_substring subs 3) in
    let typ =
      match type_string with
	  "JPEG" -> "image/jpeg"
	| "GIF" -> "image/gif"
	| "PNG" -> "image/png"
	| _ -> raise (Invalid_argument "image_identify") in
    typ, width, height
  with
      Not_found ->
	 raise (Invalid_argument "image_identify")

(* Make a thumbnail of an image.  This uses the ImageMagick program 'convert'.
 *)
let image_thumbnail data max_width max_height =
  let filename = output_tempfile data in
  let cmd = sprintf "convert -size %dx%d %s -resize %dx%d -"
	      max_width max_height filename max_width max_height in
  let in_chan = Unix.open_process_in cmd in
  let thumbnail = input_all in_chan in
  let status = Unix.close_process_in in_chan in
  unlink filename;
  (match status with
       Unix.WEXITED 0 ->		(* Convert was OK with it ... *)
	 ()
     | Unix.WEXITED n ->		(* Convert failed. *)
	 failwith ("convert: fail with error code " ^ string_of_int n)
     | Unix.WSIGNALED n ->
	 failwith ("convert: killed by signal " ^ string_of_int n)
     | Unix.WSTOPPED n ->
	 failwith ("convert: stopped by signal " ^ string_of_int n));
  let mime_type, width, height = image_identify thumbnail in
  thumbnail, mime_type, width, height

(*----- Files and MIME types. -----*)

let ws_re = Pcre.regexp "\\s+"
let ext_re = Pcre.regexp "\\.([a-z0-9]+)$"

let mime_types =
  try
    let chan = open_in "/etc/mime.types" in
    let lines = input_all_lines chan in
    close_in chan;
    let lines = List.filter (fun line ->
			       not (string_is_whitespace line) &&
			       String.length line > 0 &&
			       line.[0] <> '#') lines in
    let res = ref [] in
    List.iter (fun line ->
		 let fields = Pcre.split ~rex:ws_re line in
		 match fields with
		     [] -> assert false
		   | typ :: exts ->
		       List.iter (fun ext ->
				    res := (ext, typ) :: !res) exts) lines;
    !res
  with
      Sys_error _ -> []

let mime_type_of_filename name =
  try
    let subs = Pcre.exec ~rex:ext_re name in
    let ext = Pcre.get_substring subs 1 in
    let ext = String.lowercase ext in
    List.assoc ext mime_types
  with
      Not_found -> "application/octet-stream"