open Std
module Lid_set = Index_format.Lid_set
let { Logger.log } = Logger.for_section "index-occurrences"

let set_fname ~file (loc : Location.t) =
  let pos_fname = file in
  { loc with
    loc_start = { loc.loc_start with pos_fname };
    loc_end = { loc.loc_end with pos_fname }
  }

let decl_of_path_or_lid env namespace path lid =
  match (namespace : Shape.Sig_component_kind.t) with
  | Constructor ->
    begin match Env.find_constructor_by_name lid env with
    | exception Not_found -> None
    | { cstr_uid; cstr_loc; _ } ->
      Some { Env_lookup.uid = cstr_uid; loc = cstr_loc; namespace }
    end
  | Label ->
    begin match Env.find_label_by_name lid env with
    | exception Not_found -> None
    | { lbl_uid; lbl_loc; _ } ->
      Some { Env_lookup.uid = lbl_uid; loc = lbl_loc; namespace }
    end
  | _ -> Env_lookup.by_path path namespace env

let iterator ~current_buffer_path ~index ~reduce_for_uid =
  let add uid loc = index := Shape.Uid.Map.add_to_list uid loc !index in
  let f ~namespace env path (lid : Longident.t Location.loc) =
    log ~title:"index_buffer" "Path: %a" Logger.fmt
      (Fun.flip (Format_doc.compat Path.print) path);
    let not_ghost { Location.loc = { loc_ghost; _ }; _ } = not loc_ghost in
    let lid = { lid with loc = set_fname ~file:current_buffer_path lid.loc } in
    let index_decl () =
      begin match decl_of_path_or_lid env namespace path lid.txt with
      | (exception _) | None ->
        log ~title:"index_buffer" "Declaration not found"
      | Some decl ->
        log ~title:"index_buffer" "Found declaration: %a" Logger.fmt
          (Fun.flip Location.print_loc decl.loc);
        add decl.uid lid
      end
    in
    let reduce_and_store ~namespace lid path =
      if not_ghost lid then
        match Env.shape_of_path ~namespace env path with
        | exception Not_found -> ()
        | path_shape ->
          log ~title:"index_buffer" "Shape of path: %a" Logger.fmt
            (Fun.flip Shape.print path_shape);
          let result = reduce_for_uid env path_shape in
          begin match Locate.uid_of_result ~traverse_aliases:false result with
          | Some uid, false ->
            log ~title:"index_buffer" "Found %a (%a) wiht uid %a" Logger.fmt
              (Fun.flip Pprintast.longident lid.txt)
              Logger.fmt
              (Fun.flip Location.print_loc lid.loc)
              Logger.fmt
              (Fun.flip Shape.Uid.print uid);
            add uid lid
          | Some uid, true ->
            log ~title:"index_buffer" "Shape is approximative, found uid: %a"
              Logger.fmt
              (Fun.flip Shape.Uid.print uid);
            index_decl ()
          | None, _ ->
            log ~title:"index_buffer" "Reduction failed: missing uid";
            index_decl ()
          end
    in
    (* Shape reduction can be expensive, but the persistent memoization tables
       should make these successive reductions fast. *)
    let rec index_components namespace lid path =
      let module_ = Shape.Sig_component_kind.Module in
      match (lid.Location.txt, path) with
      | Longident.Ldot (lid', _), Path.Pdot (path', _)
      | Ldot (lid', _), Pextra_ty (Pdot (path', _), Pcstr_ty _) ->
        reduce_and_store ~namespace lid path;
        index_components module_ lid' path'
      | Lapply (lid', lid''), Papply (path', path'')
      | Lapply (lid', lid''), Pextra_ty (Papply (path', path''), Pcstr_ty _) ->
        index_components module_ lid'' path'';
        index_components module_ lid' path'
      | Longident.Lident _, _ -> reduce_and_store ~namespace lid path
      | _, _ -> ()
    in
    index_components namespace lid path
  in
  Ast_iterators.iterator_on_usages ~f

let items index (config : Mconfig.t) items =
  let module Shape_reduce = Shape_reduce.Make (struct
    let fuel = 10

    let read_unit_shape ~unit_name =
      log ~title:"read_unit_shape" "inspecting %s" unit_name;
      let cmt = Format.sprintf "%s.cmt" unit_name in
      match Cmt_cache.read (Load_path.find_normalized cmt) with
      | { cmt_impl_shape; _ } ->
        log ~title:"read_unit_shape" "shapes loaded for %s" unit_name;
        cmt_impl_shape
      | exception _ ->
        log ~title:"read_unit_shape" "failed to find %s" unit_name;
        None
  end) in
  let current_buffer_path =
    Filename.concat config.query.directory config.query.filename
  in
  let reduce_for_uid = Shape_reduce.reduce_for_uid in
  let index = ref index in
  let iterator = iterator ~current_buffer_path ~index ~reduce_for_uid in
  let () =
    match items with
    | `Impl items -> List.iter ~f:(iterator.structure_item iterator) items
    | `Intf items -> List.iter ~f:(iterator.signature_item iterator) items
  in
  !index
