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

type t = {
  index: Export_index.t;
  value_matcher: Fuzzy_path.t;
  type_matcher: Fuzzy_path.t;
}

type search_result = {
  name: string;
  source: Export_index.source;
  kind: Export_index.kind;
}
[@@deriving show]

type search_results = {
  results: (search_result * int) list;
  is_incomplete: bool;
}
[@@deriving show]

type search_options = Fuzzy_path.options = {
  first_match_can_be_weak: bool;
  num_threads: int;
  max_results: int;
  weighted: bool;
}

let default_options : search_options =
  Fuzzy_path.{ default_options with first_match_can_be_weak = false }

type candidates = {
  values: (string * int) list;
  types: (string * int) list;
}

let summarize_exports exports =
  Export_index.ExportMap.fold
    (fun (_file, kind) num (has_value, has_type, max_count) ->
      let max_count = max num max_count in
      match kind with
      | Export_index.Default -> (true, has_type, max_count)
      | Export_index.Named -> (true, has_type, max_count)
      | Export_index.NamedType -> (has_value, true, max_count)
      | Export_index.Namespace -> (true, has_type, max_count))
    exports
    (false, false, 0)

let partition_candidates index =
  Export_index.fold_names
    ~f:(fun { values; types } name exports ->
      let (has_value, has_type, max_count) = summarize_exports exports in
      let values =
        if has_value then
          (name, max_count) :: values
        else
          values
      in
      let types =
        if has_type then
          (name, max_count) :: types
        else
          types
      in
      { values; types })
    ~init:{ values = []; types = [] }
    index

let init index =
  let { values; types } = partition_candidates index in
  let value_matcher = Fuzzy_path.init values in
  let type_matcher = Fuzzy_path.init types in
  { index; value_matcher; type_matcher }

let subtract old_index { index; value_matcher; type_matcher } =
  let (index, dead_candidates) = Export_index.subtract old_index index in
  let value_matcher = Fuzzy_path.remove_candidates value_matcher dead_candidates in
  let type_matcher = Fuzzy_path.remove_candidates type_matcher dead_candidates in
  { index; value_matcher; type_matcher }

let subtract_count removed_imports { index; value_matcher; type_matcher } =
  let index = Export_index.subtract_count removed_imports index in
  { index; value_matcher; type_matcher }

let merge new_index { index; value_matcher; type_matcher } =
  let index = Export_index.merge new_index index in
  let { values; types } = partition_candidates new_index in
  let value_matcher = Fuzzy_path.add_candidates value_matcher values in
  let type_matcher = Fuzzy_path.add_candidates type_matcher types in
  { index; value_matcher; type_matcher }

(*Merge_import *)
let merge_export_import new_index { index; value_matcher; type_matcher } =
  let index = Export_index.merge_export_import new_index index in
  { index; value_matcher; type_matcher }

type query =
  | Value of string
  | Type of string

let search_result_of_export ~query name source kind =
  let open Export_index in
  match (query, kind) with
  | (Value _, (Default | Named | Namespace))
  | (Type _, NamedType) ->
    Some { name; source; kind }
  | (Value _, NamedType)
  | (Type _, (Default | Named | Namespace)) ->
    None

let score_of_export kind count fuzzy_score =
  let open Export_index in
  let type_score =
    match kind with
    | Default -> 2
    | Named
    | NamedType ->
      1
    | Namespace -> 0
  in
  count + int_of_float ((fuzzy_score ** 4.0) *. 1000.0) + type_score

(** [take ~n:20 ~index matches] will return up to 20 search results,
    where each match in [matches] might contribute multiple results.
    sets [is_incomplete] if [n] is exceeded. *)
let take ~weighted ~n ~index ~query fuzzy_matches =
  let rev_all =
    Base.List.fold fuzzy_matches ~init:[] ~f:(fun acc { Fuzzy_path.value; score = fuzzy_score } ->
        Export_index.ExportMap.fold
          (fun (source, kind) count acc ->
            match search_result_of_export ~query value source kind with
            | Some result ->
              let score =
                if weighted then
                  score_of_export kind count fuzzy_score
                else
                  int_of_float (fuzzy_score *. 1000.0)
              in
              (result, score) :: acc
            | None -> acc)
          (Export_index.find value index)
          acc
    )
  in
  let sorted =
    (* sorts the highest scores to the front *)
    Base.List.stable_sort rev_all ~compare:(fun (a, score_a) (b, score_b) ->
        match Int.compare score_b score_a with
        | 0 ->
          (match String.compare a.name b.name with
          | 0 ->
            (match Export_index.compare_export (a.source, a.kind) (b.source, b.kind) with
            | 0 ->
              (* since rev_all is reversed, [a; b] where a and b have equal
                 scores should return [b; a]. stable_sort wants to leave them
                 in the original order. *)
              1
            | k -> k)
          | k -> k)
        | k -> k
    )
  in
  let results = Base.List.take sorted n in
  let is_incomplete = Base.List.length sorted > n in
  { results; is_incomplete }

let search ?(options = default_options) query { index; value_matcher; type_matcher } =
  let (matcher, query_txt) =
    match query with
    | Value txt -> (value_matcher, txt)
    | Type txt -> (type_matcher, txt)
  in

  let max_results = options.Fuzzy_path.max_results in
  let options =
    (* if max_results is set, then increase it by 1 so that we can tell the difference
       between getting exactly max_results results (is_incomplete = false) and getting
       too many (is_incomplete = true). *)
    if max_results < max_int then
      Fuzzy_path.{ options with max_results = max_results + 1 }
    else
      options
  in

  let weighted = options.Fuzzy_path.weighted in

  Fuzzy_path.search ~options query_txt matcher |> take ~weighted ~n:max_results ~index ~query

let search_values ?options query t = search ?options (Value query) t

let search_types ?options query t = search ?options (Type query) t

let get name { index; value_matcher = _; type_matcher = _ } = Export_index.find name index

let get_values name t =
  get name t
  |> Export_index.ExportMap.filter (fun (_file_key, kind) _num -> Export_index.kind_is_value kind)

let get_types name t =
  get name t
  |> Export_index.ExportMap.filter (fun (_file_key, kind) _num -> Export_index.kind_is_type kind)

let pp fmt { index; _ } = Export_index.pp fmt index

let show { index; _ } = Export_index.show index
