(* Compiles a specified mime type file into ML code specifying      *)
(* default mime types                                               *)
(* by Travis Bemann                                                 *)
(*                                                                  *)
(* This program 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 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 Lesser General Public License for more details.              *)

open Printf

let hash_start_size = 1024

let strip_init_dot str =
  if String.length str > 0 then
    if String.get str 0 = '.' then
      Str.string_after str 1
    else str
  else ""

let rec parse_file stream hash file =
  try
    let line = input_line stream in
    let line_real =
      Strutil.strip
	begin try
	  let off = String.index line '#' in
	  Str.string_before line off
	with Not_found -> line
	end
    in
    if line_real <> "" then
      let parts = Str.split (Str.regexp "[ \t]+") line_real in
      try
	let suffix = strip_init_dot (List.nth parts 0)
	and mimetype = List.nth parts 1 in
	begin try let _ = Hashtbl.find hash suffix in ()
	with Not_found -> Hashtbl.add hash suffix mimetype
	end;
	parse_file stream hash file
      with Failure _ ->
	eprintf "mimetype_compile: %s: unable to parse line \"%s\"\n"
	  file line;
	exit 2
    else parse_file stream hash file
  with End_of_file -> ()

let dump_hash hash =
  let items = ref [] in
  Hashtbl.iter
    (fun suffix mimetype -> items := ((suffix, mimetype) :: !items))
    hash;
  List.rev !items

let load_files files use_stdin =
  let hash = Hashtbl.create hash_start_size in
  let rec step files  =
    match files with
      file :: rest ->
	begin try
	  let stream = open_in file in
	  parse_file stream hash file;
	  close_in stream;
	  step rest
	with Sys_error _ ->
	  eprintf "mimetype_compile: %s: unable to open or read" file;
	  exit 2
	end
    | [] -> ()
  in
  step files;
  if use_stdin then
    parse_file stdin hash "<stdin>"
  else ();
  dump_hash hash

let pp_mappings buf mappings =
  let rec step mappings first =
    match mappings with
      (suffix, mimetype) :: [] ->
	if first then
	  Buffer.add_string buf "  [(\""
	else
	  Buffer.add_string buf "   (\"";
	Buffer.add_string buf suffix;
	Buffer.add_string buf "\", \"";
	Buffer.add_string buf mimetype;
	Buffer.add_string buf "\")]\n"
    | (suffix, mimetype) :: rest ->
	if first then
	  Buffer.add_string buf "  [(\""
	else
	  Buffer.add_string buf "   (\"";
	Buffer.add_string buf suffix;
	Buffer.add_string buf "\", \"";
	Buffer.add_string buf mimetype;
	Buffer.add_string buf "\");\n";
	step rest false
    | [] ->
	if first then
	  Buffer.add_string buf "  []\n"
	else assert false
  in step mappings true

let generate_code mappings =
  let buf = Buffer.create 0 in
  Buffer.add_string buf
    ("(* This is a file generated with mimetype_compile *)\n"^
     "(* Do not edit this file directly; instead, modify the "^
     "mappings file *)\n(* and rerun mimetype_compile (with "^
     "make) to change this file's *)\n(* contents. *)\n\n");
  Buffer.add_string buf "let mappings =\n";
  pp_mappings buf mappings;
  Buffer.contents buf

let save_code out_file code =
  match out_file with
    Some file ->
      begin try
	let stream = open_out file in
	output_string stream code;
	close_out stream
      with Sys_error _ ->
	eprintf "mimetype_compile: %s: unable to open or write" file;
	exit 2
      end
  | None -> print_string code

let main () =
  let files = ref []
  and use_stdin = ref false
  and out_file = ref None in
  let args =
    [(Arg2.Both ('s', "stdin"), [Arg2.Set use_stdin], [],
      "Specify standard input to be used as an input file");
     (Arg2.Both_arg ('o', "output", "file"),
      [Arg2.String (fun file -> out_file := Some file)], [],
      ("Specify the generated OCaml code to be written in a file "^
       "instead of standard output"))]
  and usage_info = "mimetype_compile [options] [file ...]"
  and descr = "Use one or more files and or standard input "^
    "to generate an OCaml file encoding default MIME type mappings"
  and notes = "by Travis Bemann"
  and prefix_error = "mimetype_compile: " in
  try
    Arg2.parse args (fun file -> files := file :: !files)
      usage_info descr notes prefix_error;
    if !use_stdin || (!files <> []) then
      let mappings = load_files !files !use_stdin in
      save_code !out_file (generate_code mappings)
    else Arg2.usage args usage_info descr notes
  with Arg2.Parse_halt -> ()

let _ = main ()