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 configured declaration =
  match List.filter is_marker declaration.ptype_attributes with
  | [] -> false
  | [(_, PStr [{ pstr_desc = Pstr_eval
      ({ pexp_desc = Pexp_ident { txt = Longident.Lident "generate"; _ }; _ }, _); _ }])] -> true
  | _ -> 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 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 source = value ~loc "__fg_source" in
  let source_pattern = Pat.constraint_ ~loc (variable ~loc "__fg_source")
      (Typ.constr ~loc (ident ~loc declaration.ptype_name.txt)
         (List.map (fun _ -> Typ.any ~loc ()) declaration.ptype_params)) in
  List.concat (List.map2 (fun field short_name ->
    let name = field.pld_name.txt in
    let getter = lambda ~loc Nolabel source_pattern
        (Exp.field ~loc source (ident ~loc name)) in
    let replace expression = Exp.record ~loc [ident ~loc name, expression]
        (if List.length fields = 1 then None else Some source) in
    let setter = lambda ~loc Nolabel (variable ~loc "__fg_value")
        (lambda ~loc Nolabel source_pattern (replace (value ~loc "__fg_value"))) in
    let updater = lambda ~loc (Labelled "f") (variable ~loc "__fg_function")
        (lambda ~loc Nolabel source_pattern
           (replace (Exp.apply ~loc (value ~loc "__fg_function")
              [Nolabel, Exp.field ~loc source (ident ~loc name)]))) in
    List.map (fun (name, body) ->
      Str.value ~loc Nonrecursive [Vb.mk ~loc (variable ~loc name) body])
      [short_name, getter; "set_" ^ short_name, setter; "update_" ^ short_name, updater;
       "_" ^ short_name, Exp.tuple ~loc [getter; setter]]) 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 generated = List.concat (List.map (fun declaration ->
          if configured declaration then generate declaration else []) declarations) in
        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 () = Ast_mapper.register "fieldglass"
    (fun _ -> { Ast_mapper.default_mapper with structure })