(*
  Copyright (c) 2005, Eric Breck
  All rights reserved.

  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are met:

  1 Redistributions of source code must retain the above copyright notice, this
    list of conditions and the following disclaimer.

  2 Redistributions in binary form must reproduce the above copyright notice,
    this list of conditions and the following disclaimer in the documentation
    and/or other materials provided with the distribution.

  3 The name of the author may not be used to endorse or promote products
    derived from this software without specific prior written permission.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *)

(** stores the symbol types created so far (used as a set of strings) *)
let symbol_types : (string, bool) Hashtbl.t = Hashtbl.create 10

let mapl f l = List.fold_right (fun x l -> f x @ l) l []

let remove_anonymous_fields fields = 
  List.fold_left (fun acc -> function (_,None,_,_,_) -> acc 
    | (l,Some n,m,t,p) -> (l,n,m,t,p)::acc) [] fields

(** generate a function to print out an options record *)
let gen_print loc kind fields =
  let fields = remove_anonymous_fields fields in
  let prints = List.map (fun (loc,name,mutable_,type_,pvl) -> 
    let prtype = try 
      <:expr< $List.assoc "str" pvl$ field>>
    with Not_found ->
      match type_ with
      | <:ctyp< $lid:s$ >> when Hashtbl.mem symbol_types s ->
        <:expr< $lid:"string_of_"^s$ field >>
      | <:ctyp< int    >> -> <:expr<string_of_int field>>
      | <:ctyp< float  >> -> <:expr<string_of_float field>>
      | <:ctyp< string >> -> <:expr<field>>
      | <:ctyp< unit   >> -> <:expr<"()">>
      | <:ctyp< bool   >> -> <:expr<string_of_bool field>>
      | _ -> <:expr<"(user-specified data)">> in
    let prtype = if List.mem_assoc "default" pvl then
      <:expr<let field = record.$lid:name$ in $prtype$>>
    else
      <:expr<match record.$lid:name$ with [ None -> "None"
      | Some field -> $prtype$ ] >> in
    <:expr< output_string channel ( $str:name ^ " = "$ ^ $prtype$ ^ "\\\\n") >>
  ) fields in
  <:expr<fun channel record ->  do { $list:prints$ }  >>

let key_name s = function
  `Short -> Printf.sprintf "-%c" s.[0]
| `Long2 -> "--" ^ s
| `Long1 -> "-" ^ s

(** generate a key (e.g. -x or -k) for a record field *)
let gen_keys kind (loc,name,mutable_,type_,pvl) = 
  try [ List.assoc "key" pvl ]
  with Not_found -> 
    let name = match name with Some n -> n | None -> invalid_arg "gen_keys" in
    let quote s = <:expr<$str:String.escaped s$>> in
    List.map (fun kind -> quote (key_name name kind)) kind

(** generate field reference declaration, e.g.
    alpha:int generates let alpha = ref default-value-of-alpha 
    where the default value either user-specified or None *)
let gen_ref (loc,name,mutable_,type_,pvl) =
  let defval = try List.assoc "default" pvl with Not_found -> <:expr<None>> in
  <:patt<$lid:name$>>,<:expr<ref $defval$>>

(** generate Arg.key * spec * doc from a record spec *)
let gen_specs kind ((loc,name,mutable_,type_,pvl) as field) =
   let getname () = match name with Some name -> name 
     | None -> invalid_arg "gen_spec" in
   (* look for "spec" param, otherwise choose spec based on type of the field *)
   let v,s,spec = 
     try 
       match List.assoc "spec" pvl with 
         <:expr< "Set"      >> -> Some "True",None,"Unit"
       | <:expr< "Clear"    >> -> Some "False",None,"Unit"
       | <:expr< $str:spec$ >> -> None,None,spec
       | _ -> invalid_arg "spec"
     with Not_found -> 
       match type_ with
       | <:ctyp<$lid:s$>> when Hashtbl.mem symbol_types s ->None,Some s,"Symbol"
       | <:ctyp< int   >> -> None,None,"Int"
       | <:ctyp< float >> -> None,None,"Float"
       | <:ctyp< unit  >> -> None,None,"Unit"
       | <:ctyp< bool  >> -> None,None,"Bool"
       | _ -> None,None,"String" in

   (* look for a "help" param, otherwise use spec + name of field *)
   let help = try List.assoc "help" pvl 
              with Not_found -> <:expr<$str:spec ^ " " ^ (getname ())$>> in

   (* look for "arg" param, otherwise just assign the field to the argument of
      the command-line option.  Special-case if the user specified the Set or
      Clear specs. Also special-cased for symbol types created with this
      extension. *)
   let arg = 
     try List.assoc "arg" pvl
     with Not_found -> begin
       let arg = match v,s,spec with 
       | Some v,_,_ -> <:expr<$uid:v$>>       (* Set | Clear -> True | False *)
       | _,Some s,_ -> <:expr<$lid:s^"_of_string"$ v>> (* user symbol type *)
       | _ -> <:expr< v >> in
       let arg = if List.mem_assoc "default" pvl then arg 
                 else <:expr< Some $arg$ >> in
       let arg = <:expr< fun v -> $lid:getname ()$.val := $arg$ >> in
       match s with Some s -> <:expr<($lid:s^"_values"$,$arg$)>> | None ->  arg
     end in
   let keys = gen_keys kind field in
   List.map (fun key -> <:expr< ($key$,Arg.$uid:spec$ $arg$,$help$) >>) keys

(** generate a function to print a usage message based on a record spec *)
let gen_usage loc kind fields =
  let named_fields = remove_anonymous_fields fields in
  <:expr<fun usage ->
    let $list:List.map gen_ref named_fields$ in
    Arg.usage $Pa_o.mklistexp loc None (mapl (gen_specs kind) fields)$ 
    (Sys.argv.(0) ^ usage)>>

(** generate a function to parse command-line args based on a record spec *)
let gen_parse loc kind fields =
  let named_fields = remove_anonymous_fields fields in
  let builds = List.map (fun (loc,name,mutable_,type_,pvl) -> 
    <:patt< $lid:name$>>,<:expr<$lid:name$.val>>) named_fields in
  <:expr<fun usage ->
    let __args = ref [] in
    let $list:List.map gen_ref named_fields$ in
    do {Arg.parse $Pa_o.mklistexp loc None (mapl (gen_specs kind) fields)$
    (fun arg -> __args.val := [ arg::__args.val]) (Sys.argv.(0) ^ usage);
    ({ $list:builds$},List.rev __args.val)} >>

(** generate record type based on the param-value-annotated type.
    We need to strip param-value annotations,
    and also convert types to 'option' if no default value is provided *)
let gen_record loc id fields = 
  let fields = remove_anonymous_fields fields in
  let fields = List.map (fun (l,n,m,t,pvl) -> 
    l,n,m,(if List.mem_assoc "default" pvl then t else <:ctyp< option $t$ >>)
  ) fields in
  <:str_item<type $lid:id$ = { $list:fields$ }>>

(** generate conversion functions and a list of values for a symbol type *)
let gen_symbol_values loc id cdl =
  let str_of_sym = List.map (fun (_,n,s) -> <:patt< $uid:n$>>,None,
    (match s with None -> <:expr< $str:String.lowercase n$>>
     | Some s -> <:expr< $str:s$ >>)
  ) cdl in
  let sym_of_str = List.map (fun (_,n,s) -> 
    (match s with None -> <:patt< $str:String.lowercase n$>>
     | Some s -> <:patt< $str:s$ >>),
    None,<:expr< $uid:n$>>) cdl in
  let sym_of_str = List.rev ((<:patt< _ >>,None,<:expr<invalid_arg
    $str:id^"_of_string"$>>) :: sym_of_str) in
  let vals = List.map (fun (_,n,s) -> 
    match s with None -> <:expr<$str:String.lowercase n$>>
     | Some s -> <:expr<$str:s$>> ) cdl in
  [<:patt< $lid:id^"_values"$>>, <:expr<$Pa_o.mklistexp loc None vals$>>;
   <:patt< $lid:"string_of_"^id$>>, <:expr<fun [ $list:str_of_sym$ ]>>;
   <:patt< $lid:id^"_of_string"$>>, <:expr<fun [ $list:sym_of_str$ ]>>]

EXTEND
  GLOBAL: Pcaml.expr;
  pv_item:
    [ [i = LIDENT; "="; e = Pcaml.expr LEVEL "expr1"; ";" -> (i,e) ] ];
  pv_label_declaration:
    [ [ i = LIDENT; ":";  l = LIST0 pv_item;  t = Pcaml.ctyp LEVEL "simple"
    -> (loc, Some i, false, t,l)
      | "_"; ":";  l = LIST0 pv_item;  t = Pcaml.ctyp LEVEL "simple"
    -> (loc, None, false, t,l)
    ]];
  pv_label_declarations:
    [ [ ld = pv_label_declaration; ";"; ldl = SELF -> (ld :: ldl)
      | ld = pv_label_declaration; ";" -> [ld]
      | ld = pv_label_declaration -> [ld] ] ];
  unary_constructor_declaration:
    [ [ ci = UIDENT -> (loc, ci, None) 
       | ci = UIDENT; s = STRING -> (loc, ci, Some s)
      ] ];
  Pcaml.str_item: [
    (*
    (* declare an option type *)
    [ "type"; LIDENT "option";  
      (* kind = OPT 
       ["["; l = LIST1 [ LIDENT "long1" -> `Long1
                                  | LIDENT "short" -> `Short 
                                  | LIDENT "long2" -> `Long2 ] SEP ","; "]" -> 
                l ] ; *)
       id=LIDENT; "="; "{"; fields=pv_label_declarations; "}"->
         let kind = None in
        let kind = match kind with None -> [`Short] | Some x -> x in
        let funcs = [<:patt< $lid:"usage_" ^ id$>>, gen_usage loc kind fields;
                     <:patt< $lid:"parse_" ^ id$>>, gen_parse loc kind fields;
                     <:patt< $lid:"print_" ^ id$>>, gen_print loc kind fields] in
        (* declare type plus three functions *)
        <:str_item<declare $gen_record loc id fields$; value $list:funcs$; end>>
    (* declare a symbol type *)
    | "type"; LIDENT "option"; 
    id=LIDENT; "="; 
      cdl = LIST1 unary_constructor_declaration SEP "|" -> 
        Hashtbl.add symbol_types id true;
        let regular_cdl = List.map (fun (l,c,_) -> l,c,[]) cdl in
        <:str_item<declare type $lid:id$ = [ $list:regular_cdl$ ]; 
          value $list:gen_symbol_values loc id cdl$; end >>
          *)
    (* declare an option type *)
    [ "type"; LIDENT "option";  
       kind = OPT 
       ["["; l = LIST1 [ LIDENT "long1" -> `Long1
                                  | LIDENT "short" -> `Short 
                                  | LIDENT "long2" -> `Long2 ] SEP ","; "]" -> 
                l ] ;
        id=LIDENT; "=";
      item = [
        "{"; fields=pv_label_declarations; "}"->
          fun kind id ->
        let kind = match kind with None -> [`Short] | Some x -> x in
        let funcs = [<:patt< $lid:"usage_" ^ id$>>, gen_usage loc kind fields;
                     <:patt< $lid:"parse_" ^ id$>>, gen_parse loc kind fields;
                     <:patt< $lid:"print_" ^ id$>>, gen_print loc kind fields] in
        (* declare type plus three functions *)
        <:str_item<declare $gen_record loc id fields$; value $list:funcs$; end>>
    (* declare a symbol type *)
    | 
      cdl = LIST1 unary_constructor_declaration SEP "|" -> 
        fun kind id -> 
        Hashtbl.add symbol_types id true;
        let regular_cdl = List.map (fun (l,c,_) -> l,c,[]) cdl in
        <:str_item<declare type $lid:id$ = [ $list:regular_cdl$ ]; 
          value $list:gen_symbol_values loc id cdl$; end >>
    ] -> item kind id]
  ];
END;;