open Import
open Ast_helper
open Printf

let prefix_of_record lds =
  common_prefix (List.map lds ~f:(fun ld -> ld.pld_name.txt))

module Gen(Fixed_loc : sig val fixed_loc : bool end) = struct
  open Fixed_loc

  let gen_combinator_for_constructor ~wrapper:(wpath, wprefix, has_attrs, has_loc_stack) path ~prefix cd =
    match cd.pcd_args with
    | Pcstr_record _ ->
      (* TODO. *)
      failwith "Pcstr_record not supported"
    | Pcstr_tuple cd_args ->
      let args =
        List.mapi cd_args ~f:(fun i _ -> sprintf "x%d" i)
      in
      let exp =
        Exp.construct (Loc.mk (fqn_longident path cd.pcd_name.txt))
          (match args with
           | []  -> None
           | [x] -> Some (evar x)
           | _   -> Some (Exp.tuple (List.map args ~f:evar)))
      in
      let body =
        let fields =
          [ ( Loc.mk (fqn_longident' wpath (wprefix ^ "loc"))
            , evar "loc"
            )
          ; ( Loc.mk (fqn_longident' wpath (wprefix ^ "desc"))
            , exp
            )
          ]
        in
        let fields =
          if has_attrs then
            ( Loc.mk (fqn_longident' wpath (wprefix ^ "attributes"))
            , M.expr "[]"
            )
            :: fields
          else
            fields
        in
        let fields =
          if has_loc_stack then
            ( Loc.mk (fqn_longident' wpath (wprefix ^ "loc_stack"))
            , M.expr "[]"
            )
            :: fields
          else
            fields
        in
        Exp.record fields None
      in
      let body =
  (*      match args with
        | [] -> [%expr fun () -> [%e body]]
        | _ ->*)
          List.fold_right args ~init:body ~f:(fun arg acc ->
            M.expr "fun %a -> %a" A.patt (pvar arg) A.expr acc)
      in
  (*    let body =
        if not has_attrs then
          body
        else
          [%expr fun ?(attrs=[]) -> [%e body]]
      in*)
      let body =
        if fixed_loc then
          body
        else
          M.expr "fun ~loc -> %a" A.expr body
      in
      M.stri "let %a = %a"
        A.patt (pvar (function_name_of_id ~prefix cd.pcd_name.txt))
        A.expr body
  ;;

  let gen_combinator_for_record path ~prefix lds =
    let fields = List.map lds ~f:(fun ld -> fqn_longident path ld.pld_name.txt) in
    let funcs =
      List.map lds ~f:(fun ld ->
        map_keyword (without_prefix ~prefix ld.pld_name.txt))
    in
    let body =
      Exp.record
        (List.map2 fields funcs ~f:(fun field func ->
           (Loc.mk field, if func = "attributes" then M.expr "[]" else evar func)))
        None
    in
    let body =
      let l = List.filter funcs ~f:(fun f -> f <> "loc" && f <> "attributes") in
      match l with
      | [x] -> Exp.fun_ Nolabel None (pvar x) body
      | _ ->
        List.fold_right l ~init:body ~f:(fun func acc ->
          Exp.fun_ (Labelled func) None (pvar func) acc
        )
    in
(*    let body =
      if List.mem "attributes" ~set:funcs then
        [%expr fun ?(attrs=[]) -> [%e body]]
      else
        body
    in*)
    let body =
      if List.mem "loc" ~set:funcs && not fixed_loc then
        M.expr "fun ~loc -> %a" A.expr body
      else
        body
    in
    M.stri "let %a = %a"
      A.patt (pvar (function_name_of_path path))
      A.expr body
  ;;

  let gen_td ?wrapper path td =
    if is_loc path then
      []
    else
      match td.ptype_kind with
      | Ptype_variant cds ->
        begin match wrapper with
        | None -> []
        | Some wrapper ->
          let prefix =
            common_prefix (List.map cds ~f:(fun cd -> cd.pcd_name.txt))
          in
          List.map cds ~f:(fun cd ->
            gen_combinator_for_constructor ~wrapper path ~prefix cd)
        end
      | Ptype_record lds ->
        let prefix = prefix_of_record lds in
        [gen_combinator_for_record path ~prefix lds]
      | Ptype_abstract | Ptype_open -> []
  ;;
end

let filter_labels ~prefix lds =
  List.filter lds ~f:(fun ld ->
    match without_prefix ~prefix ld.pld_name.txt with
    | "loc" | "loc_stack" | "attributes" -> false
    | _ -> true)
;;

let is_abstract td =
  match td.ptype_kind with
  | Ptype_abstract -> true
  | _ -> false
;;

let dump fn ~ext printer x =
  let oc = open_out (fn ^ ext) in
  let ppf = Format.formatter_of_out_channel oc in
  Format.fprintf ppf "%a@." printer x;
  close_out oc

let generate filename =
  (*  let fn = Misc.find_in_path_uncap !Config.load_path (unit ^ ".cmi") in*)
  let types = get_types ~filename in
  let types_with_wrapped =
    List.map types ~f:(fun (path, td) ->
      match td.ptype_kind with
      | Ptype_record lds ->
        let prefix = prefix_of_record lds in
        let lds' = filter_labels ~prefix lds in
        (match is_wrapper ~prefix lds' with
         | None -> (path, td, None)
         | Some p ->
           let has_attrs = List.exists lds ~f:(fun ld ->
             ld.pld_name.txt = prefix ^ "attributes")
           in
           let has_loc_stack = List.exists lds ~f:(fun ld ->
             ld.pld_name.txt = prefix ^ "loc_stack")
           in
           (path, td, Some (prefix, has_attrs, has_loc_stack, p.txt)))
      | _ -> (path, td, None))
  in
  let wrapped =
    List.filter_map types_with_wrapped ~f:(fun (_, _, x) ->
      match x with
      | None -> None
      | Some (_, _, _, p) -> Some p)
  in
  let types =
    List.filter types_with_wrapped ~f:(fun (path, _, _) ->
      not (List.mem path ~set:wrapped))
    |> List.map ~f:(fun (path, td, wrapped) ->
      match wrapped with
      | None -> (path, td, None)
      | Some (prefix, has_attrs, has_loc_stack, p) ->
        (path, td, Some (prefix, has_attrs, has_loc_stack, p, List.assoc p types)))
  in
  (*  let all_types = List.map fst types in*)
  let types =
    List.sort types ~cmp:(fun (a, _, _) (b, _, _) ->
      compare a b)
  in
  let items fixed_loc =
    let module G = Gen(struct let fixed_loc = fixed_loc end) in
    List.map types ~f:(fun (path, td, wrapped) ->
      if is_abstract td then
        []
      else
        match wrapped with
        | None -> G.gen_td path td
        | Some (prefix, has_attrs, has_loc_stack, path', td') ->
          G.gen_td ~wrapper:(path, prefix, has_attrs, has_loc_stack) path' td'
    )
    |> List.flatten
  in
  let st =
    [ Str.open_ (Opn.mk (Mod.ident (Loc.lident "Import")))
    ; Str.module_ (Mb.mk (Loc.mk "M") (Mod.structure (items false)))
    ; Str.module_ (Mb.mk (Loc.mk "Make")
                     (Mod.functor_ (Loc.mk "Loc") (Some (Mty.signature [
                        Sig.value (Val.mk (Loc.mk "loc") (M.ctyp "Location.t"))
                      ]))
                        (Mod.structure
                           (M.stri "let loc = Loc.loc"
                            :: items true))))
    ]
  in
  dump "ast_builder_generated" Pprintast.structure st ~ext:".ml"

let args =
  [ ]

let usage = Printf.sprintf "%s [options] <.ml files>\n" Sys.argv.(0)

let () =
  let fns = ref [] in
  Arg.parse (Arg.align args) (fun fn -> fns := fn :: !fns) usage;
  try
    List.iter (List.rev !fns) ~f:generate
  with exn ->
    Errors.report_error Format.err_formatter exn;
    exit 2
