(*

 * optParse - Functions for parsing command line arguments.

 * Copyright (C) 2004 Bardur Arantsson

 *

 * Heavily influenced by the optparse.py module from the Python

 * standard library, but with lots of adaptation to the 'Ocaml Way'

 *

 *

 * 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 2.1 of the License, or (at your option) any later version,

 * with the special exception on linking described in file LICENSE.

 *

 * This library 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.

 *

 * You should have received a copy of the GNU Lesser General Public

 * License along with this library; if not, write to the Free Software

 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

 *)

open Printf

open ExtString

open ExtList

let terminal_width =

  try 

    int_of_string (Sys.getenv "COLUMNS")    (* Might as well use it if it's there... *)

  with

    Failure _ -> 80

  | Not_found -> 80

module GetOpt =

  struct

    type action = string -> string list -> unit

    type long_opt = string * int * action

    type short_opt = char * int * action

    exception Error of (string * string)

    let split1 haystack needle =

      try 

        let (h, x) = String.split haystack needle in h, [x] 

      with

        Invalid_string -> haystack, []

    let find_opt format_name options s =

      let rec loop l =

        match l with

          (x, y, z) :: t -> if x = s then x, y, z else loop t

        | [] -> raise (Error (format_name s, "no such option"))

      in

      loop options

    let find_short_opt options = find_opt (fun c -> sprintf "-%c" c) options

    let find_long_opt options = find_opt (fun s -> "--" ^ s) options

    let parse other find_short_opt find_long_opt args =

      let rec loop args =

        let rec gather_args name n args =

          try 

            List.split_nth n args 

          with

            List.Invalid_index _ ->

              raise (Error (name, "missing required arguments"))

        in

        let gather_long_opt s args =

          let (h, t) = split1 s "=" in

          let (_, nargs, action) = find_long_opt (String.slice ~first:2 h) in

          let (accum, args') = gather_args h (nargs - List.length t) args in

          action h (t @ accum); args'

        in

        let rec gather_short_opt_concat seen_args s k args =

          if k < String.length s then

            let ostr = sprintf "-%c" s.[k]

            and (_, nargs, action) = find_short_opt s.[k] in

            if nargs = 0 then

              begin

                action ostr [];

                gather_short_opt_concat seen_args s (k + 1) args

              end

            else if not seen_args then

              let (accum, args') = gather_args ostr nargs args in

              action ostr accum; gather_short_opt_concat true s (k + 1) args'

            else

              raise

                (Error

                   (sprintf "-%c" s.[k],

                    sprintf "option list '%s' already contains an option requiring an argument"

                      s))

          else args

        in

        let gather_short_opt s k args =

          let ostr = sprintf "-%c" s.[k] in

          let (_, nargs, action) = find_short_opt s.[k] in

          if nargs = 0 then gather_short_opt_concat false s k args

          else

            let (accum, args') =

              let h = String.slice ~first:(k+1) s in

              if String.length h = 0 then gather_args ostr nargs args

              else

                let (t, args'') = gather_args ostr (nargs - 1) args in

                h :: t, args''

            in

            action ostr accum; args'

        in

        match args with

          [] -> []

        | arg :: args' ->

            if arg = "--" then args'

            else if String.starts_with arg "--" then

              loop (gather_long_opt arg args')

            else if arg = "-" then begin other arg; loop args' end

            else if String.starts_with arg "-" then

              loop (gather_short_opt arg 1 args')

            else begin other arg; loop args' end

      in

      let args' = loop args in List.iter other args'

  end

module Opt =

  struct

    exception No_value

    exception Option_error of string * string

    exception Option_help

    type 'a t = { 

      option_set : string -> string list -> unit;

      option_set_value : 'a -> unit;

      option_get : unit -> 'a option;

      option_metavars : string list;

      option_defhelp : string option 

    }

    let get opt =

      match opt.option_get () with

        Some x -> x

      | None -> raise No_value

    let set opt v =

      opt.option_set_value v

    let is_set opt = Option.is_some (opt.option_get ())

    let opt opt = opt.option_get ()

    let value_option metavar default coerce errfmt =

      let data = ref default in

      {

        option_metavars = [metavar]; 

        option_defhelp = None;

        option_get = (fun _ -> !data);

        option_set_value = (fun x -> data := Some x);

        option_set =

         (fun option args ->

            let arg = List.hd args in

              try 

                data := Some (coerce arg)

              with

                  exn -> raise (Option_error (option, errfmt exn arg)))

      }

    let callback_option metavar coerce errfmt f =

      {

        option_metavars = [metavar]; 

        option_defhelp = None;

        option_get = (fun _ -> Some ());

        option_set_value = (fun () -> ());

        option_set =

         (fun option args ->

            let arg = List.hd args in

            let datum = ref None in

              begin 

              try 

                datum := Some (coerce arg)

              with

                  exn -> raise (Option_error (option, errfmt exn arg))

              end;

              Option.may f !datum)

      }

  end

module StdOpt =

  struct

    open Opt

    let store_const ?default const =

      let data = ref default in

      {

        option_metavars = []; 

        option_defhelp = None;

        option_get = (fun _ -> !data);

        option_set_value = (fun x -> data := Some x);

        option_set = fun _ _ -> data := Some const

      }

    let store_true () = store_const ~default:false true

    let store_false () = store_const ~default:true false

    let int_option ?default ?(metavar = "INT") () =

      value_option metavar default int_of_string

        (fun _ s -> sprintf "invalid integer value '%s'" s)

    let int_callback ?(metavar = "INT") =

      callback_option metavar int_of_string

        (fun _ s -> sprintf "invalid integer value '%s'" s)

    let float_option ?default ?(metavar = "FLOAT") () =

      value_option metavar default float_of_string

        (fun _ s -> sprintf "invalid floating point value '%s'" s)

    let float_callback ?(metavar = "FLOAT") =

      callback_option metavar float_of_string

        (fun _ s -> sprintf "invalid floating point value '%s'" s)

    let str_option ?default ?(metavar = "STR") () =

      value_option metavar default (fun s -> s) (fun _ _ -> "cannot happen")

    let str_callback ?(metavar = "STR") =

      callback_option metavar (fun s -> s) (fun _ _ -> "cannot happen")

    let count_option ?(dest = ref 0) ?(increment = 1) () =

      {

        option_metavars = []; 

        option_defhelp = None;

        option_get = (fun _ -> Some !dest);

        option_set_value = (fun x -> dest := x);

        option_set = fun _ _ -> dest := !dest + increment

      }

    let incr_option ?(dest = ref 0) = 

      count_option ~dest ~increment:1

    let decr_option ?(dest = ref 0) = 

      count_option ~dest ~increment:(-1)

    let help_option () =

      {

        option_metavars = [];

        option_defhelp = Some "show this help message and exit";

        option_get = (fun _ -> raise No_value);

        option_set_value = (fun _ -> ());

        option_set = fun _ _ -> raise Option_help

      }

    let version_option vfunc =

      {

        option_metavars = [];

        option_defhelp = Some "show program's version and exit";

        option_get = (fun _ -> raise No_value);

        option_set_value = (fun _ -> ());

        option_set = fun _ _ -> print_endline (vfunc ()); exit 0

      }

  end

module Formatter =

  struct

    (* Note that the whitespace regexps must NOT treat the non-breaking

       space character as whitespace. *)

    let whitespace = "\t\n\013\014\r "

    let split_into_chunks s =

      let buf = Buffer.create (String.length s) in

      let flush () =

        let s = Buffer.contents buf

        in

          Buffer.clear buf;

          s

      in

      let rec loop state accum i =

        if (i

          if ((state && not (String.contains whitespace s.[i])) || 

              ((not state) && String.contains whitespace s.[i])) then

            if Buffer.length buf > 0 then

               loop (not state) (flush () :: accum) i 

             else 

               loop (not state) accum i

          else

            begin

              Buffer.add_char buf s.[i];

              loop state accum (i+1)

            end

        else

          if Buffer.length buf > 0 then

            flush () :: accum

          else 

            accum

      in

        List.rev (loop false [] 0)

    let is_whitespace s =

      let rec loop i =

        if i

          if String.contains whitespace s.[i] then

            loop (i+1)

          else 

            false

        else 

          true

      in

        loop 0

    let expand_tabs ?(tab_size = 8) s =

      let len = String.length s in

      let spaces n = String.make n ' '

      and b = Buffer.create len in

      let rec expand i col =

        if i < len then

          match s.[i] with

            '\t' ->

              let n = tab_size - col mod tab_size in

              Buffer.add_string b (spaces n);

              expand (i + 1) (col + n)

          | '\n' -> 

              Buffer.add_string b "\n";

              expand (i + 1) 0

          | c -> 

              Buffer.add_char b c;

              expand  (i + 1) (col + 1)

      in

      expand 0 0; 

      Buffer.contents b

    let wrap ?(initial_indent = 0) ?(subsequent_indent = 0) text _width =

      let wrap_chunks_line width acc =

        let rec wrap (chunks, cur_line, cur_len) =

          match chunks with

            [] -> [], cur_line, cur_len

          | hd :: tl ->

              let l = String.length hd in

              if cur_len + l <= width then

                wrap (tl, hd :: cur_line, cur_len + l)

              else chunks, cur_line, cur_len

        in

        wrap acc

      in

      let wrap_long_last_word width (chunks, cur_line, cur_len) =

        match chunks with

          [] -> [], cur_line, cur_len

        | hd :: tl ->

            let l = String.length hd in

            if l > width then

              match cur_line with

                [] -> tl, [hd], cur_len + l

              | _ -> chunks, cur_line, cur_len

            else chunks, cur_line, cur_len

      in

      let wrap_remove_last_ws (chunks, cur_line, cur_len) =

        match cur_line with

          [] -> chunks, cur_line, cur_len

        | hd :: tl ->

            if is_whitespace hd then chunks, tl, cur_len - String.length hd

            else chunks, cur_line, cur_len

      in

      let rec wrap_chunks_lines chunks lines =

        let indent =

          match lines with

            [] -> initial_indent

          | _ -> subsequent_indent

        in

        let width = _width - indent in

        match chunks with

          hd :: tl ->

            if is_whitespace hd && lines <> [] then wrap_chunks_lines tl lines

            else (* skip *)

              let (chunks', cur_line, _) =

                wrap_remove_last_ws

                  (wrap_long_last_word width

                     (wrap_chunks_line width (chunks, [], 0)))

              in

              wrap_chunks_lines chunks'

                ((String.make indent ' ' ^

                    String.concat "" (List.rev cur_line)) ::

                   lines)

        | [] -> List.rev lines

      in

      let chunks = split_into_chunks (expand_tabs text) in

      wrap_chunks_lines chunks []

    let fill ?(initial_indent = 0) ?(subsequent_indent = 0) text width =

      String.concat "\n" (wrap ~initial_indent ~subsequent_indent text width)

    type t = { 

      indent : unit -> unit;

      dedent : unit -> unit;

      format_usage : string -> string;

      format_heading : string -> string;

      format_description : string -> string;

      format_option : char list * string list -> string list -> 

                                             string option -> string

    }

    let format_option_strings short_first (snames, lnames) metavars =

      let metavar = String.concat " " metavars in

      let lopts =

        List.map

          (match metavar with

             "" -> (fun z -> sprintf "--%s" z)

           | _ -> fun z -> sprintf "--%s=%s" z metavar)

          lnames

      and sopts = List.map (fun x -> sprintf "-%c%s" x metavar) snames in

      match short_first with

        true -> String.concat ", " (sopts @ lopts)

      | false -> String.concat ", " (lopts @ sopts)

    let indented_formatter ?level:(extlevel = ref 0)

      ?indent:(extindent = ref 0) ?(indent_increment = 2) 

      ?(max_help_position = 24) ?(width = terminal_width - 1) 

      ?(short_first = true) () =

      let indent = ref 0

      and level = ref 0 in

      let help_position = ref max_help_position

      and help_width = ref (width - max_help_position) in

      {

        indent =

         (fun () ->

            indent := !indent + indent_increment;

            level := !level + 1;

            extindent := !indent;

            extlevel := !level);

        dedent =

         (fun () ->

            indent := !indent - indent_increment;

            level := !level - 1;

            assert (!level >= 0);

            extindent := !indent;

            extlevel := !level);

        format_usage = (fun usage -> sprintf "usage: %s\n" usage);

        format_heading =

         (fun heading -> sprintf "%*s%s:\n\n" !indent "" heading);

        format_description =

         (fun description ->

            let x =

              fill ~initial_indent:(!indent) ~subsequent_indent:(!indent)

                description (width - !indent)

            in

              if not (String.ends_with x "\n") then x ^ "\n\n" else x ^ "\n");

        format_option =

         fun names metavars help ->

           let opt_width = !help_position - !indent - 2 in

           let opt_strings =

             format_option_strings short_first names metavars

           in

           let buf = Buffer.create 256 in

           let indent_first =

             if String.length opt_strings > opt_width then

               begin

                 bprintf buf "%*s%s\n" !indent "" opt_strings; !help_position

               end

             else

               begin

                 bprintf buf "%*s%-*s  " !indent "" opt_width opt_strings; 0

               end

           in

           Option.may

             (fun option_help ->

                let lines = wrap option_help !help_width in

                match lines with

                  h :: t ->

                    bprintf buf "%*s%s\n" indent_first "" h;

                    List.iter

                      (fun x -> bprintf buf "%*s%s\n" !help_position "" x) t

                | [] -> ())

             help;

           let contents =

             Buffer.contents buf

           in

             if String.length contents > 0 && not (String.ends_with contents "\n") then

               contents ^ "\n"

             else

               contents

      }

    let titled_formatter ?(level = ref 0) ?(indent = ref 0) 

      ?(indent_increment = 0) ?(max_help_position = 24) 

      ?(width = terminal_width - 1) ?(short_first = true) 

      () =

      let formatter =

        indented_formatter ~level ~indent ~indent_increment ~max_help_position

          ~width ~short_first ()

      in

      let format_heading h =

        let c =

          match !level with

            0 -> '='

          | 1 -> '-'

          | _ -> failwith "titled_formatter: Too much indentation"

        in

        sprintf "%*s%s\n%*s%s\n\n" !indent "" (String.capitalize h) !indent ""

          (String.make (String.length h) c)

      in

      let format_usage usage =

        sprintf "%s  %s\n" (format_heading "Usage") usage

      in

      { formatter with 

          format_usage = format_usage;

          format_heading = format_heading

      }

  end

open Opt

open Formatter

module OptParser =

  struct

    exception Option_conflict of string

    type group = { 

      og_heading : string;

      og_description : string option;

      og_options :

        ((char list * string list) * string list * string option) RefList.t;

      og_children : group RefList.t 

    }

    type t = { 

      op_usage : string;

      op_status : int;

      op_suppress_usage : bool;

      op_prog : string;

      op_formatter : Formatter.t;

      op_long_options : GetOpt.long_opt RefList.t;

      op_short_options : GetOpt.short_opt RefList.t;

      op_groups : group 

    }

    let unprogify optparser s =

      (snd (String.replace ~str:s ~sub:"%prog" ~by:optparser.op_prog))

    let add optparser ?(group = optparser.op_groups) ?help ?(hide = false)

      ?short_name ?(short_names = []) ?long_name ?(long_names = []) opt =

      let lnames =

        match long_name with

            None -> long_names

          | Some x -> x :: long_names

      and snames =

        match short_name with

            None -> short_names

          | Some x -> x :: short_names

      in

      if lnames = [] && snames = [] then

        failwith "Options must have at least one name"

      else

        (* Checking for duplicates: *)

        let snames' =

          List.fold_left (fun r (x, _, _) -> x :: r) []

            (RefList.to_list optparser.op_short_options)

        and lnames' =

          List.fold_left (fun r (x, _, _) -> x :: r) []

            (RefList.to_list optparser.op_long_options)

        in

        let sconf =

          List.filter (fun e -> List.exists (( = ) e) snames') snames

        and lconf =

          List.filter (fun e -> List.exists (( = ) e) lnames') lnames

        in

        if List.length sconf > 0 then

          raise (Option_conflict (sprintf "-%c" (List.hd sconf)))

        else if List.length lconf > 0 then

          raise (Option_conflict (sprintf "--%s" (List.hd lconf)));

        (* Add to display list. *)

        if not hide then

          RefList.add group.og_options

            ((snames, lnames), opt.option_metavars,

             (match help with

                  None -> opt.option_defhelp

                | Some _ -> help));

        (* Getopt: *)

        let nargs = List.length opt.option_metavars in

          List.iter

            (fun short ->

               RefList.add optparser.op_short_options

               (short, nargs, opt.option_set))

            snames;

          List.iter

            (fun long ->

               RefList.add optparser.op_long_options

               (long, nargs, opt.option_set))

            lnames

    let add_group optparser ?(parent = optparser.op_groups) ?description heading =

      let g =

        {

          og_heading = heading; 

          og_description = description;

          og_options = RefList.empty (); 

          og_children = RefList.empty ()

        }

      in

      RefList.add parent.og_children g; g

    let make ?(usage = "%prog [options]") ?(status = 1) ?description ?version

      ?(suppress_usage = false) ?(suppress_help = false) ?prog 

      ?(formatter = Formatter.indented_formatter ()) () =

      let optparser =

        {

          op_usage = usage; 

          op_status = status;

          op_suppress_usage = suppress_usage;

          op_prog = Option.default (Filename.basename Sys.argv.(0)) prog;

          op_formatter = formatter; 

          op_short_options = RefList.empty ();

          op_long_options = RefList.empty ();

          op_groups = {

            og_heading = "options"; 

            og_options = RefList.empty ();

            og_children = RefList.empty ();

            og_description = description

          }

        }

      in

      Option.may                         (* Add version option? *)

        (fun version ->

           add optparser ~long_name:"version"

             (StdOpt.version_option

                (fun () -> unprogify optparser version)))

        version;

      if not suppress_help then              (* Add help option? *)

        add optparser ~short_name:'h' ~long_name:"help"

          (StdOpt.help_option ());

      optparser

    let format_usage optparser eol =

      match optparser.op_suppress_usage with

        true -> ""

      | false ->

          unprogify optparser

            (optparser.op_formatter.format_usage optparser.op_usage) ^ eol

    let error optparser ?(chn = stderr) ?status message =

      fprintf chn "%s%s: %s\n" (format_usage optparser "\n") optparser.op_prog

        message;

      flush chn;

      exit (Option.default optparser.op_status status)

    let usage optparser ?(chn = stdout) () =

      let rec loop g =

        (* Heading: *)

        output_string chn

          (optparser.op_formatter.format_heading g.og_heading);

        optparser.op_formatter.indent ();

        (* Description: *)

        Option.may

          (fun x ->

             output_string chn (optparser.op_formatter.format_description x))

          g.og_description;

        (* Options: *)

        RefList.iter

          (fun (names, metavars, help) ->

             output_string chn

               (optparser.op_formatter.format_option names metavars help))

          g.og_options;

        (* Child groups: *)

        output_string chn "\n";

        RefList.iter loop g.og_children;

        optparser.op_formatter.dedent ()

      in

      output_string chn (format_usage optparser "\n");

      loop optparser.op_groups;

      flush chn

    let parse optparser ?(first = 0) ?last argv =

      let args = RefList.empty ()

      and n =

        match last with

          None -> Array.length argv - first

        | Some m -> m - first + 1

      in

      begin 

        try

          GetOpt.parse (RefList.push args)

            (GetOpt.find_short_opt

               (RefList.to_list optparser.op_short_options))

            (GetOpt.find_long_opt (RefList.to_list optparser.op_long_options))

            (Array.to_list (Array.sub argv first n))

        with

            GetOpt.Error (opt, errmsg) ->

              error optparser (sprintf "option '%s': %s" opt errmsg)

          | Option_error (opt, errmsg) ->

              error optparser (sprintf "option '%s': %s" opt errmsg)

          | Option_help -> usage optparser (); exit 0

      end;

      List.rev (RefList.to_list args)

    let parse_argv optparser = 

      parse optparser ~first:1 Sys.argv

  end