open Std
open Typedtree

let {Logger. log} = Logger.for_section "construct"

type values_scope = Null | Local
type what = Modtype | Mod

exception Not_allowed of string
exception Not_a_hole
exception Modtype_not_found of what * string
exception No_constraint

let () =
  Location.register_error_of_exn (function
    | Not_a_hole ->
      Some (Location.error "Construct only works on holes.")
    | Modtype_not_found (Modtype, s) ->
      let txt = Format.sprintf "Module type not found: %s" s in
      Some (Location.error txt)
    | Modtype_not_found (Mod, s) ->
      let txt = Format.sprintf "Module not found: %s" s in
      Some (Location.error txt)
    | No_constraint ->
      Some (Location.error
        "Could not find a module type to construct from. \
        Check that you used a correct constraint.")
    | _ -> None
  )
module Util = struct
  open Misc_utils.Path
  open Types

  let predef_types =
    let tbl = Hashtbl.create 14 in
    let () =
      let constant c =
        Ast_helper.Exp.constant c
      in
      let construct s =
        Ast_helper.Exp.construct (Location.mknoloc (Longident.Lident s)) None
      in
      let ident s =
        Ast_helper.Exp.ident (Location.mknoloc (Longident.Lident s))
      in
      List.iter ~f:(fun (k, v) -> Hashtbl.add tbl k v)
        Parsetree.[
          Predef.path_int, constant (Pconst_integer("0", None)) ;
          Predef.path_float, constant (Pconst_float("0.0", None)) ;
          Predef.path_char, constant (Pconst_char 'c') ;
          Predef.path_string,
            constant (Pconst_string("", Location.none, None)) ;
          Predef.path_bool, construct "false" ;
          Predef.path_unit, construct "()" ;
          Predef.path_exn, ident "exn" ;
          Predef.path_array, Ast_helper.Exp.array [] ;
          Predef.path_nativeint, constant (Pconst_integer("0", Some 'n')) ;
          Predef.path_int32, constant (Pconst_integer("0", Some 'l')) ;
          Predef.path_int64, constant (Pconst_integer("0", Some 'L')) ;
          Predef.path_lazy_t, Ast_helper.Exp.lazy_ (construct "()")
        ]
    in
    tbl

  let prefix env ~env_check path name =
    to_shortest_lid ~env ~env_check ~name path

  let maybe_prefix env ~env_check path name =
    match Warnings.is_active (Disambiguated_name "") with
    | false -> Longident.Lident name
    | true -> prefix env ~env_check path name

  let var_of_id id = Location.mknoloc @@ Ident.name id

  let type_to_string t =
    Printtyp.type_expr (Format.str_formatter) t;
    Format.flush_str_formatter ()

  let unifiable env type_expr type_expected =
    let snap = Btype.snapshot () in
    try
      Ctype.unify env type_expected type_expr |> ignore;
      Some snap
    with Ctype.Unify _  ->
      (* Unification failure *)
      Btype.backtrack snap;
      None

  let typeable env exp type_expected =
    let snap = Btype.snapshot () in
    let typeable =
      match Typecore.type_expect env exp (Typecore.mk_expected type_expected) with
      | (_ : Typedtree.expression) -> true
      | exception _ -> false
    in
    if not typeable then
      log ~title:"constructor" "%a does not have the expected type %a"
        Logger.fmt (fun fmt -> Printast.expression 0 fmt exp)
        Logger.fmt (fun fmt -> Printtyp.type_expr fmt type_expected);
    Btype.backtrack snap;
    typeable

  let is_in_stdlib path =
    Path.head path |> Ident.name = "Stdlib"

  (** [find_values_for_type env typ] searches the environment [env] for
  {i values} with a return type compatible with [typ] *)
  let find_values_for_type env typ =
    let aux name path value_description acc =
      (* [check_type| checks return type compatibility and lists parameters *)
      let rec check_type type_expr params =
        let type_expr = Transient_expr.repr type_expr in
        (* TODO is this test general enough ? *)
        match unifiable env (Transient_expr.type_expr type_expr) typ with
        | Some snap ->
          (* This will be called multiple times so we need to backtrack
              See c-simple, test 6.2b for an example *)
          Btype.backtrack snap;
          Some params
        | None ->
          begin match type_expr.desc with
          | Tarrow (arg_label, _, te, _) -> check_type te (arg_label::params)
          | _ -> None
      end
      in
      (* TODO we should probably sort the results better *)
      match is_in_stdlib path, check_type value_description.val_type [] with
      | false, Some params ->
        Path.Map.add path (name, value_description, params) acc
      | _, _ -> acc
    in
    (* We look for values in the current scope and in local unonpend submodules.
      We also exclude the Stdlib modules from the search. *)
    let fold_values path acc = Env.fold_values aux path env acc in
    let init = fold_values None Path.Map.empty in
    Env.fold_modules (fun _name path _module_decl acc ->
      if not (is_in_stdlib path) && not (is_opened env path) then
        (* We ignore opened modules. That means that is a value of an opened
          module has been shadowed we won't suggest the one in the opened
          module. *)
        fold_values (Some (Untypeast.lident_of_path path)) acc
      else acc) None env init


  (** The idents_table is used to keep track of already used names when
  generating function arguments in the same expression *)
  let idents_table ~keywords =
    let table = Hashtbl.create 50 in
    (* We add keywords to the table so they are always numbered *)
    List.iter keywords ~f:(fun k -> Hashtbl.add table k (-1));
    table

  (* Given a list [l] of n elements which are lists of choices,
    [combination l] is a list of all possible combinations of
    these choices (cartesian product). For example:

    let l = [["a";"b"];["1";"2"]; ["x"]];;
    combinations l;;
    - : string list list =
    [["a"; "1"; "x"]; ["b"; "1"; "x"];
     ["a"; "2"; "x"]; ["b"; "2"; "x"]]

    If the input is the empty list, the result is
    the empty list singleton list.
    *)
  let combinations l =
    List.fold_left l
      ~init:[[]]
      ~f:(fun acc_l choices_for_arg_i ->
        List.fold_left choices_for_arg_i
          ~init:[]
          ~f:(fun acc choice_arg_i ->
            let choices = List.map acc_l
              ~f:(fun l -> List.rev (choice_arg_i :: l))
            in
            List.rev_append acc choices))

  (** [panache2 l1 l2] returns a new list containing an interleaving of the
  values in [l1] and [l2] *)
  let panache2 l1 l2 =
    let rec aux acc l1 l2 =
      match l1, l2 with
      | [], [] -> List.rev acc
      | tl, [] | [], tl -> List.rev_append acc tl
      | a::tl1, b::tl2 -> aux (a::b::acc) tl1 tl2
  in aux [] l1 l2

  (* Given a list [l] of n lists, [panache l] flattens the list
    by starting with the first element of each, then the second one etc. *)
  let panache l =
    List.fold_left ~init:[] ~f:panache2 l
end

module Gen = struct
  open Types

  (* [make_value] generates the PAST repr of a value applied to holes *)
  let make_value env (path, (name, _value_description, params)) =
    let open Ast_helper in
    let env_check = Env.find_value_by_name in
    let lid = Location.mknoloc (Util.prefix env ~env_check path name) in
    let params = List.map params
      ~f:(fun label -> label, Exp.hole ())
    in
    if List.length params > 0 then
      Exp.(apply (ident lid) params)
    else Exp.ident lid

  (* We never perform deep search when constructing modules *)
  let rec module_ env =
    let open Ast_helper in function
    | Mty_ident path -> begin
      try
        let m = Env.find_modtype path env in
        match m.mtd_type with
        | Some t -> module_ env t
        | None -> raise Not_found
      with Not_found ->
        let name = Ident.name (Path.head path) in
        raise (Modtype_not_found (Modtype, name))
      end
    | Mty_signature sig_items ->
      let env = Env.add_signature sig_items env in
      Mod.structure @@ structure env sig_items
    | Mty_functor (param, out) ->
      let param = match param with
        | Unit -> Parsetree.Unit
        | Named (id, in_) ->
          Parsetree.Named (
            Location.mknoloc (Option.map ~f:Ident.name id),
            Ptyp_of_type.module_type in_)
      in
      Mod.functor_ param @@ module_ env out
    | Mty_alias path ->
      begin try let m = Env.find_module path env in
        module_ env m.md_type
        with Not_found ->
          let name = Ident.name (Path.head path) in
          raise (Modtype_not_found (Mod, name))
      end
    | Mty_for_hole -> Mod.hole ()
  and structure_item env =
    let open Ast_helper in
    function
    | Sig_value (id, _vd, _visibility) ->
      let vb = Vb.mk (Pat.var (Util.var_of_id id)) (Exp.hole ()) in
      Str.value Nonrecursive [ vb ]
    | Sig_type (id, type_declaration, rec_flag, _visibility) ->
      let td = Ptyp_of_type.type_declaration id type_declaration in
      let rec_flag = match rec_flag with
        | Trec_first | Trec_next -> Asttypes.Recursive
        | Trec_not -> Nonrecursive
      in (* mutually recursive types are really handled by [structure] *)
      Str.type_ rec_flag [td]
    | Sig_modtype (id, { mtd_type; _ }, _visibility) ->
      let mtd = Ast_helper.Mtd.mk
        ?typ:(Option.map ~f:Ptyp_of_type.module_type mtd_type)
        @@ Util.var_of_id id
      in
      Ast_helper.Str.modtype mtd
    | Sig_module (id, _, mod_decl, _, _) ->
      let module_binding =
        Ast_helper.Mb.mk
          (Location.mknoloc (Some (Ident.name id)))
          @@ module_ env mod_decl.md_type
      in
      Str.module_ module_binding
    | Sig_typext (id, ext_constructor, _, _) ->
      let lid =
        Untypeast.lident_of_path ext_constructor.ext_type_path
        |> Location.mknoloc
      in
      Str.type_extension @@ Ast_helper.Te.mk
        ~attrs:ext_constructor.ext_attributes
        ~params:[]
        ~priv:ext_constructor.ext_private
        lid
        [Ptyp_of_type.extension_constructor id ext_constructor]
    | Sig_class_type (id, _class_type_decl, _, _) ->
      let str = Format.asprintf "Construct does not handle class types yet. \
      Please replace this comment by [%s]'s definition." (Ident.name id) in
      Str.text [ Docstrings.docstring str Location.none ] |> List.hd
    | Sig_class (id, _class_decl, _, _) ->
      let str = Format.asprintf "Construct does not handle classes yet. \
      Please replace this comment by [%s]'s definition." (Ident.name id) in
      Str.text [ Docstrings.docstring str Location.none ] |> List.hd
  and structure env (items : Types.signature_item list) =
    List.map (Ptyp_of_type.group_items items) ~f:(function
      | Ptyp_of_type.Item item -> structure_item env item
      | Ptyp_of_type.Type (rec_flag, type_decls) ->
        Ast_helper.Str.type_ rec_flag type_decls)

  (* [expression values_scope ~depth env ty] generates a list of PAST
    expressions that could fill a hole of type [ty] in the environment [env].
    [depth] regulates the deep construction of recursive values. If
    [values_scope] is set to [Local] the returned list will also contains
    local values to choose from *)
  let rec expression ~idents_table values_scope ~depth =
    let exp_or_hole env typ =
      if depth > 1 then
        (* If max_depth has not been reached we recurse *)
        expression ~idents_table values_scope ~depth:(depth - 1) env typ
      else
        (* else we return a hole *)
        [ Ast_helper.Exp.hole () ]
    in
    let arrow_rhs env typ =
      match (Transient_expr.repr typ).desc with
      | Tarrow _ -> expression ~idents_table values_scope ~depth env typ
      | _ -> exp_or_hole env typ
    in

    (* [make_arg] tries to provide a nice default name for function args *)
    let make_arg =
      let make_i n i =
        Hashtbl.replace idents_table n i;
        Printf.sprintf "%s_%i" n i
      in
      let uniq_name env n =
        let id = Ident.create_local n in
        try
          let i = Hashtbl.find idents_table n + 1 in
          make_i n i
        with Not_found ->
          try
            let _ = Env.find_value (Path.Pident id) env in
            make_i n 0
          with Not_found -> Hashtbl.add idents_table n 0; n
      in
      fun env label ty ->
        let open Asttypes in
        match label with
        | Labelled s | Optional s ->
            (* Pun for labelled arguments *)
            Ast_helper.Pat.var ( Location.mknoloc s), s
        | Nolabel -> begin match get_desc ty with
          | Tconstr (path, _, _) ->
            let name = uniq_name env (Path.last path) in
            Ast_helper.Pat.var (Location.mknoloc name), name
          | _ -> Ast_helper.Pat.any (), "_" end
    in

    let constructor env type_expr path constrs =
      log ~title:"constructors" "[%s]"
        (String.concat ~sep:"; "
          (List.map constrs ~f:(fun c -> c.Types.cstr_name)));
      (* [make_constr] builds the PAST repr of a type constructor applied
      to holes *)
      let make_constr env path type_expr cstr_descr =
        let ty_args, ty_res, _ = Ctype.instance_constructor cstr_descr in
        match Util.unifiable env type_expr ty_res with
        | Some snap ->
          let lid =
            Util.maybe_prefix env
              ~env_check:Env.find_constructor_by_name
              path cstr_descr.cstr_name
            |> Location.mknoloc
          in
          let args = List.map ty_args ~f:(exp_or_hole env) in
          let args_combinations = Util.combinations args in
          let exps = List.map args_combinations
            ~f:(function
              | [] -> None
              | [e] -> Some e
              | l -> Some (Ast_helper.Exp.tuple l))
          in
          Btype.backtrack snap;
          List.filter_map exps
            ~f:(fun exp ->
              let exp = Ast_helper.Exp.construct lid exp in
              (* For gadts not all combinations will be valid.
                See Test 6.1b in c-simple.t for an example.

                We therefore check that constructed expressions
                can be typed. *)
              if Util.typeable env exp type_expr
              then Some exp else (
                log ~title:"constructor" "%s's type is not unifiable with %a"
                  cstr_descr.Types.cstr_name
                  Logger.fmt (fun fmt -> Printtyp.type_expr fmt type_expr);
                None))
        | None -> []
      in
      List.map constrs ~f:(make_constr env path type_expr)
      (* [constrs] are ordered inversly to a source code declaration.
         We reverse it to match it and provide better UX *)
      |> List.rev
      |> Util.panache
    in

    let variant env _typ row_desc =
      let fields =
        List.filter
          ~f:(fun (_lbl, row_field) -> match row_field_repr row_field with
            | Rpresent _
            | Reither (true, [], _)
            | Reither (false, [_], _) -> true
            | _ -> false)
          (row_fields row_desc)
        (* [row_fields] are ordered inversly to a source code declaration.
           We reverse it to match it and provide better UX *)
        |> List.rev
      in
      match fields with
      | [] -> raise (Not_allowed "empty variant type")
      | row_descrs ->
        List.map row_descrs ~f:(fun (lbl, row_field) ->
          (match row_field_repr row_field with
            | Reither (false, [ty], _) | Rpresent (Some ty) ->
              List.map ~f:(fun s -> Some s) (exp_or_hole env ty)
            | _ -> [None])
            |> List.map ~f:(fun e ->
              Ast_helper.Exp.variant lbl e)
            )
        |> List.flatten
        |> List.rev
    in

    let record env typ path labels =
      log ~title:"record labels" "[%s]"
        (String.concat ~sep:"; "
          (List.map labels ~f:(fun l -> l.Types.lbl_name)));

      let labels = List.map labels ~f:(fun ({ lbl_name; _ } as lbl) ->
        let _, arg, res = Ctype.instance_label true lbl in
        Ctype.unify env res typ ;
        let lid =
          Util.maybe_prefix env
            ~env_check:Env.find_label_by_name
            path lbl_name
          |> Location.mknoloc
        in
        let exprs = exp_or_hole env arg in
        lid, exprs)
      in

      let lbl_lids, lbl_exprs = List.split labels in
      Util.combinations lbl_exprs
      |> List.map
          ~f:(fun lbl_exprs ->
            let labels = List.map2 lbl_lids lbl_exprs
              ~f:(fun lid exp -> (lid, exp))
            in
            Ast_helper.Exp.record labels None)
    in

    (* Given a typed hole, there is two possible forms of constructions:
      - Use the type's definition to propose the correct type constructors,
      - Look for values in the environment with compatible return type. *)
    fun env typ ->
      log ~title:"construct expr" "Looking for expressions of type %s"
        (Util.type_to_string typ);
      let rtyp =
        Ctype.full_expand ~may_forget_scope:true env typ
      in
      let constructed_from_type = match get_desc rtyp with
        | Tlink _ | Tsubst _ ->
          assert false
        | Tpoly (texp, _)  ->
          (* We are not going "deeper" so we don't call [exp_or_hole] here *)
          expression ~idents_table values_scope ~depth env texp
        | Tunivar _ | Tvar _ ->
          [ ]
        | Tconstr (path, [texp], _) when path = Predef.path_lazy_t ->
          (* Special case for lazy *)
          let exps = exp_or_hole env texp in
          List.map exps ~f:Ast_helper.Exp.lazy_
        | Tconstr (path, _params, _) ->
          (* If this is a "basic" type we propose a default value *)
          begin try
            [ Hashtbl.find Util.predef_types path ]
          with Not_found ->
            let def = Env.find_type_descrs path env in
            match def with
            | Type_variant (constrs, _) -> constructor env rtyp path constrs
            | Type_record (labels, _) -> record env rtyp path labels
            | Type_abstract | Type_open -> []
          end
        | Tarrow (label, tyleft, tyright, _) ->
          let argument, name = make_arg env label tyleft in
          let value_description = {
              val_type = tyleft;
              val_kind = Val_reg;
              val_loc = Location.none;
              val_attributes = [];
              val_uid = Uid.mk ~current_unit:(Env.get_unit_name ());
            }
          in
          let env = Env.add_value (Ident.create_local name) value_description env in
          let exps = arrow_rhs env tyright in
          List.map exps ~f:(Ast_helper.Exp.fun_ label None argument)
        | Ttuple types ->
          let choices = List.map types ~f:(exp_or_hole env)
            |> Util.combinations
          in
          List.map choices  ~f:Ast_helper.Exp.tuple
        | Tvariant row_desc -> variant env rtyp row_desc
        | Tpackage (path, lids_args) -> begin
          let open Ast_helper in
          try
            let ty =
              Typemod.modtype_of_package env Location.none path lids_args
            in
            let ast =
              Exp.constraint_
                (Exp.pack (module_ env ty))
                (Ptyp_of_type.core_type typ)
            in
            [ ast ]
          with Typemod.Error _ ->
            let name = Ident.name (Path.head path) in
            raise (Modtype_not_found (Modtype, name)) end
        | Tobject (fields, _) ->
          let rec aux acc fields =
            match get_desc fields with
            | Tnil -> acc
            | Tvar _ | Tunivar _ -> acc
            | Tfield ("*dummy method*", _, _, fields) -> aux acc fields
            | Tfield (name, _, type_expr, fields) ->
              let exprs = exp_or_hole env type_expr
                |> List.map ~f:(fun expr ->
                  let open Ast_helper in
                  Cf.method_ (Location.mknoloc name) Asttypes.Public
                  @@ Ast_helper.Cf.concrete Asttypes.Fresh expr)
              in
              aux (exprs :: acc) fields
            | _ ->
              failwith @@ Format.asprintf
                "Unexpected type constructor in fields list: %a"
                Printtyp.type_expr fields
          in
          let all_fields = aux [] fields |> Util.combinations in
          List.map all_fields ~f:(fun fields ->
            let open Ast_helper in
            Exp.object_ @@ Ast_helper.Cstr.mk (Pat.any ()) fields)
        | Tfield _ | Tnil -> failwith "Found a field type outside an object"
      in
      let matching_values =
        if values_scope = Local then
          Path.Map.bindings (Util.find_values_for_type env typ)
          |> List.map ~f:(make_value env)
        else []
      in
      List.append constructed_from_type matching_values
end

let needs_parentheses e = match e.Parsetree.pexp_desc with
  | Pexp_fun _
  | Pexp_lazy _
  | Pexp_apply _
  | Pexp_variant (_, Some _)
  | Pexp_construct (_, Some _)
  -> true
  | _ -> false

let to_string_with_parentheses exp =
  let f : _ format6 =
    if needs_parentheses exp then "(%a)"
    else "%a"
  in
  Format.asprintf f Pprintast.expression exp

let node ?(depth = 1) ~(config : Mconfig.t) ~keywords ~values_scope node =
  Warnings.with_state config.ocaml.warnings
    (fun () ->
       match node with
       | Browse_raw.Expression { exp_type; exp_env; _ } ->
         let idents_table = Util.idents_table ~keywords in
         Gen.expression ~idents_table values_scope ~depth exp_env exp_type
         |> List.map ~f:to_string_with_parentheses
       | Browse_raw.Module_expr
           { mod_desc = Tmod_constraint _ ; mod_type; mod_env; _ }
       | Browse_raw.Module_expr
           {  mod_desc = Tmod_apply _; mod_type; mod_env; _ } ->
         let m = Gen.module_ mod_env mod_type in
         [ Format.asprintf "%a" Pprintast.module_expr m ]
       | Browse_raw.Module_expr _
       | Browse_raw.Module_binding _ ->
         (* Constructible modules have an explicit constraint or are functor
            applications. In other cases we do not know what to construct.

            It is ok to raise here, since Warnings.with_state handles it. *)
         raise No_constraint
       | _ ->
         (* As above, it is ok to raise here. *)
         raise Not_a_hole)
