(**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

open Utils_js
open ServerEnv

module Result = Core_result
let (>>=) = Result.(>>=)
let (>>|) = Result.(>>|)

open FindRefsUtils

(* The default visitor does not provide all of the context we need when visiting an object key. In
 * particular, we need the location of the enclosing object literal. *)
class ['acc] object_key_visitor ~init = object(this)
  inherit ['acc, Loc.t] Flow_ast_visitor.visitor ~init as super

  method! expression (exp: (Loc.t, Loc.t) Flow_ast.Expression.t) =
    let open Flow_ast.Expression in
    begin match exp with
    | loc, Object x ->
      this#visit_object_literal loc x
    | _ -> ()
    end;
    super#expression exp

  method private visit_object_literal (loc: Loc.t) (obj: (Loc.t, Loc.t) Flow_ast.Expression.Object.t) =
    let open Flow_ast.Expression.Object in
    let get_prop_key =
      let open Property in
      function Init { key; _ } | Method { key; _ } | Get { key; _ } | Set { key; _ } -> key
    in
    let { properties; comments= _ } = obj in
    properties
    |> List.iter begin function
      | SpreadProperty _ -> ()
      | Property (_, prop) -> prop |> get_prop_key |> this#visit_object_key loc
    end

  method private visit_object_key
      (_literal_loc: Loc.t)
      (_key: (Loc.t, Loc.t) Flow_ast.Expression.Object.Property.key) =
    ()
end

module ObjectKeyAtLoc : sig
  (* Given a location, returns Some (enclosing_literal_loc, prop_loc, name) if the given location
   * points to an object literal key. The first location returned is the location for the entire
   * enclosing object literal. This is because later, we need to figure out which types are related
   * to this object literal which is easier to do when we have the location of the actual object
   * literal than if we only had the location of a single key. *)
  val get: (Loc.t, Loc.t) Flow_ast.program -> Loc.t -> (Loc.t * Loc.t * string) option
end = struct
  class object_key_finder target_loc = object(this)
    inherit [(Loc.t * Loc.t * string) option] object_key_visitor ~init:None
    method! private visit_object_key
        (literal_loc: Loc.t)
        (key: (Loc.t, Loc.t) Flow_ast.Expression.Object.Property.key) =
      let open Flow_ast.Expression.Object in
      match key with
      | Property.Identifier (prop_loc, { Flow_ast.Identifier.name; comments= _ }) when Loc.contains prop_loc target_loc ->
        this#set_acc (Some (literal_loc, prop_loc, name))
      | _ -> ()
  end

  let get ast target_loc =
    let finder = new object_key_finder target_loc in
    finder#eval finder#program ast
end

(* If the given type refers to an object literal, return the location of the object literal.
 * Otherwise return None *)
let get_object_literal_loc ty : Loc.t option =
  let open Type in
  let open Reason in
  let reason_desc =
    reason_of_t ty
    (* TODO look into unwrap *)
    |> desc_of_reason ~unwrap:false
  in
  match reason_desc with
  | RObjectLit -> Some (Type.def_loc_of_t ty |> ALoc.to_loc)
  | _ -> None

type def_kind =
  (* Use of a property, e.g. `foo.bar`. Includes type of receiver (`foo`) and name of the property
   * `bar` *)
  | Use of Type.t * string
  (* In a class, where a property/method is defined. Includes the type of the class and the name
  of the property. *)
  | Class_def of Type.t * string (* name *) * bool (* static *)
  (* In an object type. Includes the location of the property definition and its name. *)
  | Obj_def of Loc.t * string (* name *)
  (* List of types that the object literal flows into directly, as well as the name of the
   * property. *)
  | Use_in_literal of Type.t Nel.t * string (* name *)

let set_def_loc_hook prop_access_info literal_key_info target_loc =
  let set_prop_access_info new_info =
    let set_ok info = prop_access_info := Ok (Some info) in
    let set_err err = prop_access_info := Error err in
    match !prop_access_info with
      | Error _ -> ()
      | Ok None -> prop_access_info := Ok (Some new_info)
      | Ok (Some info) -> begin match info, new_info with
        | Use _, Use _
        | Class_def _, Class_def _
        | Obj_def _, Obj_def _ ->
          (* Due to generate_tests, we sometimes see hooks firing multiple times for the same
           * location. This is innocuous and we should take the last result. *)
          set_ok new_info
        (* Literals can flow into multiple types. Include them all. *)
        | Use_in_literal (types, name), Use_in_literal (new_types, new_name) ->
          if name = new_name then
            set_ok (Use_in_literal (Nel.rev_append new_types types, name))
          else
            set_err "Names did not match"
        (* We should not see mismatches. *)
        |  Use _, _ | Class_def _, _ | Obj_def _, _ | Use_in_literal _, _ ->
          set_err "Unexpected mismatch between definition kind"
      end
  in
  let use_hook ret _ctxt name loc ty =
    let loc = ALoc.to_loc loc in
    begin if Loc.contains loc target_loc then
      set_prop_access_info (Use (ty, name))
    end;
    ret
  in
  let class_def_hook _ctxt ty static name loc =
    let loc = ALoc.to_loc loc in
    if Loc.contains loc target_loc then
      set_prop_access_info (Class_def (ty, name, static))
  in
  let obj_def_hook _ctxt name loc =
    let loc = ALoc.to_loc loc in
    if Loc.contains loc target_loc then
      set_prop_access_info (Obj_def (loc, name))
  in
  let export_named_hook name loc =
    let loc = ALoc.to_loc loc in
    if Loc.contains loc target_loc then
      set_prop_access_info (Obj_def (loc, name))
  in
  let obj_to_obj_hook _ctxt obj1 obj2 =
    match get_object_literal_loc obj1, literal_key_info with
    | Some loc, Some (target_loc, _, name) when loc = target_loc ->
      let open Type in
      begin match obj2 with
      | DefT (_, _, ObjT _) ->
        set_prop_access_info (Use_in_literal (Nel.one obj2, name))
      | _ -> ()
      end
    | _ -> ()
  in

  Type_inference_hooks_js.set_member_hook (use_hook false);
  Type_inference_hooks_js.set_call_hook (use_hook ());
  Type_inference_hooks_js.set_class_member_decl_hook class_def_hook;
  Type_inference_hooks_js.set_obj_prop_decl_hook obj_def_hook;
  Type_inference_hooks_js.set_export_named_hook export_named_hook;
  Type_inference_hooks_js.set_obj_to_obj_hook obj_to_obj_hook

let unset_hooks () =
  Type_inference_hooks_js.reset_hooks ()

type single_def_info =
  | Class of Loc.t
  (* An object was found. *)
  | Object of Loc.t

(* If there are multiple relevant definition locations (e.g. the request was issued on an object
 * literal which is associated with multiple types) then there will be multiple locations in no
 * particular order. *)
type property_def_info = single_def_info Nel.t

type def_info =
  | Property of property_def_info * string (* name *)
  | CJSExport of Loc.t

let display_name_of_def_info = function
  | Property (_, name) -> name
  | CJSExport _ -> "module.exports"

let loc_of_single_def_info = function
  | Class loc -> loc
  | Object loc -> loc

let all_locs_of_property_def_info def_info =
  def_info
  |> Nel.map loc_of_single_def_info

let all_locs_of_def_info = function
  | Property (def_info, _) -> all_locs_of_property_def_info def_info
  | CJSExport loc -> Nel.one loc

type def_loc =
  (* We found a class property. Include all overridden implementations. Superclass implementations
   * are listed last. *)
  | FoundClass of Loc.t Nel.t
  (* We found an object property. *)
  | FoundObject of Loc.t
  | FoundUnion of def_loc Nel.t
  (* This means we resolved the receiver type but did not find the definition. If this happens
   * there must be a type error (which may be suppresssed) *)
  | NoDefFound
  (* This means it's a known type that we deliberately do not currently support. *)
  | UnsupportedType
  (* This means it's not well-typed, and could be anything *)
  | AnyType

let debug_string_of_locs locs =
  locs |> Nel.to_list |> Core_list.map ~f:Loc.debug_to_string |> String.concat ", "

(* Disable the unused value warning -- we want to keep this around for debugging *)
[@@@warning "-32"]
let debug_string_of_single_def_info = function
  | Class loc -> spf "Class (%s)" (Loc.debug_to_string loc)
  | Object loc -> spf "Object (%s)" (Loc.debug_to_string loc)

let debug_string_of_property_def_info def_info =
  def_info
  |> Nel.map debug_string_of_single_def_info
  |> Nel.to_list
  |> String.concat ", "
  |> spf "[%s]"

let debug_string_of_def_info = function
  | Property (def_info, name) ->
    spf "Property (%s, %s)" (debug_string_of_property_def_info def_info) name
  | CJSExport loc ->
    spf "CJSExport (%s)" (Loc.debug_to_string loc)

let rec debug_string_of_def_loc = function
  | FoundClass locs -> spf "FoundClass (%s)" (debug_string_of_locs locs)
  | FoundObject loc -> spf "FoundObject (%s)" (Loc.debug_to_string loc)
  | FoundUnion def_locs ->
    Nel.to_list def_locs |> Core_list.map ~f:debug_string_of_def_loc |> String.concat ", "
    |> spf "FoundUnion (%s)"
  | NoDefFound -> "NoDefFound"
  | UnsupportedType -> "UnsupportedType"
  | AnyType -> "AnyType"
(* Re-enable the unused value warning *)
[@@@warning "+32"]

let extract_instancet cx ty : (Type.t, string) result =
  let open Type in
  let resolved = Members.resolve_type cx ty in
  match resolved with
    | ThisClassT (_, t)
    | DefT (_, _, PolyT (_, _, ThisClassT (_, t), _)) -> Ok t
    | _ ->
      let type_string = string_of_ctor resolved in
      Error ("Expected a class type to extract an instance type from, got " ^ type_string)

(* Must be called with the result from Members.extract_type *)
let get_def_loc_from_extracted_type cx extracted_type name =
  extracted_type
  |> Members.extract_members cx
  |> Members.to_command_result
  >>| fun map -> match SMap.get name map with
    | None -> None
    (* Currently some types (e.g. spreads) do not contain locations for their properties. For now
     * we'll just treat them as if the properties do not exist, but once this is fixed this case
     * should be promoted to an error *)
    | Some (None, _) -> None
    | Some (Some loc, _) -> Some loc

let rec extract_def_loc cx ty name : (def_loc, string) result =
  let resolved = Members.resolve_type cx ty in
  extract_def_loc_resolved cx resolved name

(* The same as get_def_loc_from_extracted_type except it recursively checks for overridden
 * definitions of the member in superclasses and returns those as well *)
and extract_def_loc_from_instancet cx extracted_type super name : (def_loc, string) result =
  let current_class_def_loc = get_def_loc_from_extracted_type cx extracted_type name in
  current_class_def_loc
  >>= begin function
    | None -> Ok NoDefFound
    | Some loc ->
      extract_def_loc cx super name
      >>= begin function
        | FoundClass lst ->
            (* Avoid duplicate entries. This can happen if a class does not override a method,
             * so the definition points to the method definition in the parent class. Then we
             * look at the parent class and find the same definition. *)
            let lst =
              if Nel.hd lst = loc then
                lst
              else
                Nel.cons loc lst
            in
            Ok (FoundClass lst)
        | FoundObject _ -> Error "A superclass should be a class, not an object"
        | FoundUnion _ -> Error "A superclass should be a class, not a union"
        (* If the superclass does not have a definition for this method, or it is for some reason
         * not a class type, or we don't know its type, just return the location we already know
         * about. *)
        | NoDefFound | UnsupportedType | AnyType -> Ok (FoundClass (Nel.one loc))
      end
  end

and extract_def_loc_resolved cx ty name : (def_loc, string) result =
  let open Members in
  let open Type in
  match extract_type cx ty with
    | Success (DefT (_, _, InstanceT (_, super, _, _))) as extracted_type ->
        extract_def_loc_from_instancet cx extracted_type super name
    | Success (DefT (_, _, ObjT _)) | SuccessModule _ as extracted_type ->
        get_def_loc_from_extracted_type cx extracted_type name
        >>| begin function
          | None -> NoDefFound
          | Some loc -> FoundObject loc
        end
    | Success (UnionT (_, rep)) ->
        let union_members =
          UnionRep.members rep
          |> Core_list.map ~f:(fun member -> extract_def_loc cx member name)
          |> Result.all
        in
        union_members
        >>= begin fun members ->
          Nel.of_list members
          |> Result.of_option ~error:"Union should have at least one member"
        end
        >>| begin fun members_nel ->
          FoundUnion members_nel
        end
    | Success _
    | FailureNullishType
    | FailureUnhandledType _ ->
        Ok UnsupportedType
    | FailureAnyType ->
        Ok AnyType

(* Takes the file key where the module reference appeared, as well as the module reference, and
 * returns the file name for the module that the module reference refers to. *)
let file_key_of_module_ref ~reader file_key module_ref =
  let resolved = Module_js.find_resolved_module
    ~reader:(Abstract_state_reader.State_reader reader)
    ~audit:Expensive.warn
    file_key
    module_ref
  in
  Module_heaps.Reader.get_file ~reader ~audit:Expensive.warn resolved

let def_info_of_typecheck_results cx props_access_info =
  let def_info_of_class_member_locs locs =
    (* We want to include the immediate implementation as well as all superclass implementations.
     * If we wanted a mode where superclass implementations were not included, for example, we
     * could choose to take only the first extracted location. *)
    Nel.map (fun loc -> Class loc) locs
  in
  let def_info_of_type name ty =
    let rec def_info_of_def_loc = function
    | FoundClass locs -> Some (def_info_of_class_member_locs locs)
    | FoundObject loc -> Some (Nel.one (Object loc))
    | FoundUnion def_locs ->
      def_locs
      |> Nel.map def_info_of_def_loc
      |> Nel.cat_maybes
      |> Option.map ~f:Nel.concat
    | NoDefFound
    | UnsupportedType
    | AnyType -> None
    in
    extract_def_loc cx ty name >>| def_info_of_def_loc
  in
  match props_access_info with
    | None -> Ok None
    | Some (Obj_def (loc, name)) ->
        Ok (Some (Nel.one (Object loc), name))
    | Some (Class_def (ty, name, static)) ->
        if static then
          (* Here, `ty` ends up resolving to `ObjT` so we lose the knowledge that this is a static
           * property. This means that we don't get the fancy look-up-the-inheritance-chain behavior
           * that we get with class instances. That would be nice to add at some point. *)
          def_info_of_type name ty
          >>| Option.map ~f:(fun def_info -> (def_info, name))
        else
          (* We get the type of the class back here, so we need to extract the type of an instance *)
          extract_instancet cx ty >>= fun ty ->
          begin extract_def_loc_resolved cx ty name >>= function
            | FoundClass locs -> Ok (Some (def_info_of_class_member_locs locs, name))
            | FoundUnion _
            | FoundObject _ -> Error "Expected to extract class def info from a class"
            | _ -> Error "Unexpectedly failed to extract definition from known type"
          end
    | Some (Use (ty, name)) ->
        def_info_of_type name ty
        >>| Option.map ~f:(fun def_info -> (def_info, name))
    | Some (Use_in_literal (types, name)) ->
        let def_infos_result =
          Nel.map (def_info_of_type name) types
          |> Nel.result_all
        in
        def_infos_result >>| fun def_infos ->
        Nel.cat_maybes def_infos
        |> Option.map ~f:(Nel.concat)
        |> Option.map ~f:(fun def_info -> (def_info, name))

let add_literal_properties literal_key_info def_info =
  (* If we happen to be on an object property, include the location of that
   * property as a def loc. We don't want to do that above because:
   * (a) We could also encounter a `Use_in_literal` if this object literal flows
   * into another object type. This would force us to make props_access_info a
   * list and add additional complexity just for the sake of this one case.
   * (b) We would have to add a type inference hook, which we are trying to
   * avoid. *)
  let def_info = match def_info, literal_key_info with
  | None, None -> Ok None
  | Some _, None -> Ok def_info
  | None, Some (_, loc, name) -> Ok (Some (Nel.one (Object loc), name))
  | Some (defs, name1), Some (_, loc, name2) ->
    if name1 <> name2 then
      Error "Unexpected name mismatch"
    else
      Ok (Some (Nel.cons (Object loc) defs, name1))
  in
  Result.map def_info ~f:(Option.map ~f:(fun (prop_def_info, name) -> Property (prop_def_info, name)))

let get_def_info
    ~reader genv env profiling file_key content loc: (def_info option, string) result Lwt.t =
  let options = genv.options in
  let props_access_info = ref (Ok None) in
  compute_ast_result options file_key content %>>= fun (ast, file_sig, info) ->
  (* Check if it's an exported symbol *)
  let loc = Option.value (ImportExportSymbols.find_related_symbol file_sig loc) ~default:loc in
  let info = Docblock.set_flow_mode_for_ide_command info in
  let literal_key_info: (Loc.t * Loc.t * string) option = ObjectKeyAtLoc.get ast loc in
  let%lwt cx =
    set_def_loc_hook props_access_info literal_key_info loc;
    let%lwt cx, _ = Profiling_js.with_timer_lwt profiling ~timer:"MergeContents" ~f:(fun () ->
      let%lwt () =
        Types_js.ensure_checked_dependencies ~options ~reader ~env file_key file_sig
      in
      Lwt.return @@
        Merge_service.merge_contents_context ~reader options file_key ast info file_sig
    ) in
    Lwt.return cx
  in
  unset_hooks ();
  !props_access_info %>>= fun props_access_info ->
  let def_info = def_info_of_typecheck_results cx props_access_info in
  let def_info = def_info >>= add_literal_properties literal_key_info in
  let def_info = def_info >>= function
  | Some _ as def_info -> Ok def_info
  | None ->
    (* Check if we are on a CJS import/export. These cases are not covered above since the type
     * system hooks don't quite get us what we want. *)
    let export_loc =
      let open File_sig.With_Loc in
      List.fold_left begin fun acc -> function
      | Require { source = (_, module_ref); require_loc; _ } ->
        if Loc.contains require_loc loc then begin match acc with
        | Error _ -> acc
        | Ok (Some _) -> Error "Did not expect multiple requires to match one location"
        | Ok None ->
          let external_file_sig =
            let filename = file_key_of_module_ref ~reader file_key module_ref in
            Option.bind filename (Parsing_heaps.Reader.get_file_sig ~reader)
          in
          Result.return @@ Option.bind external_file_sig begin fun external_file_sig ->
            match external_file_sig.module_sig.module_kind with
            | CommonJS { mod_exp_loc=Some loc; _ } -> Some loc
            | _ -> None
          end
        end else acc
      | _ -> acc
      end (Ok None) file_sig.module_sig.requires
    in
    let export_loc = export_loc >>| function
    | Some _ as x -> x
    | None ->
      let open File_sig.With_Loc in
      match file_sig.module_sig.module_kind with
      | CommonJS { mod_exp_loc=Some mod_exp_loc; _ } ->
        if Loc.contains mod_exp_loc loc then Some mod_exp_loc
        else None
      | _ -> None
    in
    Result.map export_loc ~f:(Option.map ~f:(fun x -> CJSExport x))
  in
  Lwt.return @@ def_info
