open Asttypes
open Parsetree
open Ast_helper

let ident ~loc name = Location.mkloc (Longident.Lident name) loc
let value ~loc name = Exp.ident ~loc (ident ~loc name)
let variable ~loc name = Pat.var ~loc (Location.mkloc name loc)
let lambda ~loc label pattern body = Exp.fun_ ~loc label None pattern body

let is_marker (name, _) = name.txt = "fieldglass" || name.txt = "lens"

let check_identifier ~loc name =
  let input = Lexing.from_string name in
  let valid = try
      match Lexer.token input with
      | Parser.LIDENT parsed when parsed = name -> Lexer.token input = Parser.EOF
      | _ -> false
    with Lexer.Error _ -> false in
  if not valid then Location.raise_errorf ~loc
      "fieldglass: invalid generated identifier %S" name

type prefix = Plain | Named of string | Type_name
type options = {
  field_prefix : prefix;
  get_prefix : string option;
  set_prefix : string;
  update_prefix : string;
  self_first : bool;
  function_label : string option;
  getters : bool;
  setters : bool;
  updaters : bool;
  lenses : bool;
}
let defaults = { field_prefix = Plain; get_prefix = None;
                 set_prefix = "set"; update_prefix = "update";
                 self_first = false; function_label = Some "f";
                 getters = true; setters = true; updaters = true; lenses = true }

let word expression =
  match expression.pexp_desc with
  | Pexp_ident { txt = Longident.Lident name; _ }
  | Pexp_constant (Pconst_string (name, _)) -> name
  | _ -> Location.raise_errorf ~loc:expression.pexp_loc
      "fieldglass: expected an identifier or string"

let flag name expression =
  match expression.pexp_desc with
  | Pexp_ident { txt = Longident.Lident value; _ } when value = name -> true
  | Pexp_construct ({ txt = Longident.Lident "true"; _ }, None) -> true
  | Pexp_construct ({ txt = Longident.Lident "false"; _ }, None) -> false
  | _ -> Location.raise_errorf ~loc:expression.pexp_loc
      "fieldglass: expected a boolean for %s" name

let configure options (label, expression) =
  match label with
  | Labelled "field_prefix" -> { options with field_prefix = Named (word expression) }
  | Labelled "field_prefix_from_type" ->
      { options with field_prefix = if flag "field_prefix_from_type" expression then Type_name else Plain }
  | Labelled "get_prefix" -> { options with get_prefix = Some (word expression) }
  | Labelled "set_prefix" -> { options with set_prefix = word expression }
  | Labelled "update_prefix" -> { options with update_prefix = word expression }
  | Labelled "self_arg_first" -> { options with self_first = flag "self_arg_first" expression }
  | Labelled "func_named_arg" ->
      let name = word expression in
      check_identifier ~loc:expression.pexp_loc name;
      { options with function_label = Some name }
  | Labelled "func_no_named_arg" ->
      { options with function_label = if flag "func_no_named_arg" expression then None else Some "f" }
  | Labelled "no_get" -> { options with getters = not (flag "no_get" expression) }
  | Labelled "no_set" -> { options with setters = not (flag "no_set" expression) }
  | Labelled "no_update" -> { options with updaters = not (flag "no_update" expression) }
  | Labelled "no_lens" -> { options with lenses = not (flag "no_lens" expression) }
  | Labelled "just_lens" ->
      let functions = not (flag "just_lens" expression) in
      { options with getters = functions; setters = functions; updaters = functions }
  | _ -> Location.raise_errorf ~loc:expression.pexp_loc
      "fieldglass: unrecognised configuration option"

let configured declaration =
  match List.filter is_marker declaration.ptype_attributes with
  | [] -> None
  | [(_, PStr [{ pstr_desc = Pstr_eval (expression, []); _ }])] ->
      let arguments = match expression.pexp_desc with
        | Pexp_ident { txt = Longident.Lident "generate"; _ } -> []
        | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Longident.Lident "generate"; _ }; _ }, arguments) -> arguments
        | _ -> Location.raise_errorf ~loc:expression.pexp_loc
            "fieldglass: expected generate followed by labelled options" in
      let seen = Hashtbl.create 8 in
      List.iter (fun (label, argument) ->
        match label with
        | Labelled name ->
            if Hashtbl.mem seen name then Location.raise_errorf ~loc:argument.pexp_loc
                "fieldglass: duplicate option %s" name;
            Hashtbl.add seen name ()
        | _ -> Location.raise_errorf ~loc:argument.pexp_loc
            "fieldglass: expected a labelled option") arguments;
      Some arguments
  | _ -> Location.raise_errorf ~loc:declaration.ptype_loc
      "fieldglass: expected [@@fieldglass generate]"

let field_names fields =
  let names = List.map (fun field -> field.pld_name.txt) fields in
  match names with
  | [] | [_] -> names
  | first :: rest ->
      let shared = List.fold_left (fun limit name ->
        let limit = min limit (String.length name) in
        let rec compare i =
          if i < limit && first.[i] = name.[i] then compare (i + 1) else i in
        compare 0) (String.length first) rest in
      let rec boundary n =
        if n = 0 || first.[n - 1] = '_' then n else boundary (n - 1) in
      let count = boundary shared in
      List.map (fun name -> String.sub name count (String.length name - count)) names

let generate options declaration =
  let loc = declaration.ptype_loc in
  let fields = match declaration.ptype_kind with
    | Ptype_record fields -> fields
    | _ -> Location.raise_errorf ~loc "fieldglass: expected a record type"
  in
  let writes = options.setters || options.updaters || options.lenses in
  if declaration.ptype_private = Private && writes then
    Location.raise_errorf ~loc "fieldglass: private records only support getters";
  List.iter (fun field -> match field.pld_type.ptyp_desc with
    | Ptyp_poly (_ :: _, _) when writes -> Location.raise_errorf ~loc:field.pld_loc
        "fieldglass: polymorphic fields only support getters"
    | _ -> ()) fields;
  let source = value ~loc "__fg_source" in
  let record_type () = Typ.constr ~loc (ident ~loc declaration.ptype_name.txt)
      (List.map (fun _ -> Typ.any ~loc ()) declaration.ptype_params) in
  let source_pattern = Pat.constraint_ ~loc (variable ~loc "__fg_source")
      (record_type ()) in
  let arguments self_first label name body =
    let self = lambda ~loc Nolabel source_pattern in
    let other = lambda ~loc label (variable ~loc name) in
    if self_first then self (other body) else other (self body) in
  List.concat (List.map2 (fun field short_name ->
    let name = field.pld_name.txt in
    let short_name = match options.field_prefix with
      | Plain -> short_name
      | Named prefix -> prefix ^ "_" ^ short_name
      | Type_name -> declaration.ptype_name.txt ^ "_" ^ short_name in
    let getter = lambda ~loc Nolabel source_pattern
        (Exp.field ~loc source (ident ~loc name)) in
    let replace expression = Exp.constraint_ ~loc
        (Exp.record ~loc [ident ~loc name, expression]
           (if List.length fields = 1 then None else Some source)) (record_type ()) in
    let setter self_first = arguments self_first Nolabel "__fg_value"
        (replace (value ~loc "__fg_value")) in
    let label = match options.function_label with None -> Nolabel | Some name -> Labelled name in
    let updater = arguments options.self_first label "__fg_function"
           (replace (Exp.apply ~loc (value ~loc "__fg_function")
              [Nolabel, Exp.field ~loc source (ident ~loc name)])) in
    List.concat (List.map (fun (enabled, name, body) ->
      if enabled then begin
        check_identifier ~loc name;
        [Str.value ~loc Nonrecursive [Vb.mk ~loc (variable ~loc name) body]]
      end else [])
      [options.getters, (match options.get_prefix with None -> short_name | Some prefix -> prefix ^ "_" ^ short_name), getter;
       options.setters, options.set_prefix ^ "_" ^ short_name, setter options.self_first;
       options.updaters, options.update_prefix ^ "_" ^ short_name, updater;
       options.lenses, "_" ^ short_name, Exp.tuple ~loc [getter; setter false]])) fields (field_names fields))

let structure mapper items =
  List.concat (List.map (fun item ->
    let item = Ast_mapper.default_mapper.structure_item mapper item in
    match item.pstr_desc with
    | Pstr_type (recursive, declarations) ->
        let configurations = List.map configured declarations in
        let generated =
          if List.exists (function Some _ -> true | None -> false) configurations then
            let arguments = List.concat (List.map (function Some args -> args | None -> []) configurations) in
            let options = List.fold_left configure defaults arguments in
            List.concat (List.map (generate options) declarations)
          else [] in
        let names = Hashtbl.create 16 in
        List.iter (fun item -> match item.pstr_desc with
          | Pstr_value (_, [{ pvb_pat = { ppat_desc = Ppat_var name; _ }; _ }]) ->
              if Hashtbl.mem names name.txt then Location.raise_errorf ~loc:name.loc
                  "fieldglass: duplicate generated binding %s; customise the prefixes" name.txt;
              Hashtbl.add names name.txt ()
          | _ -> ()) generated;
        let declarations = List.map (fun declaration ->
          { declaration with ptype_attributes =
              List.filter (fun attribute -> not (is_marker attribute)) declaration.ptype_attributes }) declarations in
        { item with pstr_desc = Pstr_type (recursive, declarations) } :: generated
    | _ -> [item]) items)

let signature_item mapper item =
  (match item.psig_desc with
   | Psig_type (_, declarations) ->
       List.iter (fun declaration ->
         if List.exists is_marker declaration.ptype_attributes then
           Location.raise_errorf ~loc:declaration.ptype_loc
             "fieldglass: generate accessors in an implementation, not a signature") declarations
   | _ -> ());
  Ast_mapper.default_mapper.signature_item mapper item

let () = Ast_mapper.register "fieldglass"
    (fun _ -> { Ast_mapper.default_mapper with structure; signature_item })