module Id = Refs.Id
module Lang = Refs.Lang

let nm id = Id.name id

let starts_with ~prefix s =
  let lp = String.length prefix in
  String.length s >= lp && String.sub s 0 lp = prefix

(* Names containing [__] are the compiler/dune mangling for wrapped internal
   modules; odoc treats them as hidden, so we don't offer them as completions. *)
let hidden s =
  let n = String.length s in
  let rec loop i = i + 1 < n && (s.[i] = '_' && s.[i + 1] = '_' || loop (i + 1)) in
  loop 0

(* {1 Kind tags} *)

(* Recognised reference kind tags -> the canonical kind we tag children with.
   Mirrors odoc's [match_*_reference_kind] (including the deprecated aliases),
   so [value-], [modtype-] &c. still filter. *)
let kind_table =
  [
    ("module", "module");
    ("module-type", "module-type");
    ("modtype", "module-type");
    ("type", "type");
    ("val", "val");
    ("value", "val");
    ("exception", "exception");
    ("exn", "exception");
    ("constructor", "constructor");
    ("const", "constructor");
    ("field", "field");
    ("recfield", "field");
    ("extension", "extension");
    ("extension-decl", "extension-decl");
    ("class", "class");
    ("class-type", "class-type");
    ("classtype", "class-type");
    ("method", "method");
    ("instance-variable", "instance-variable");
    ("label", "label");
    ("section", "label");
    ("page", "page");
  ]

(* Split a final component into (kind-prefix-verbatim, canonical-kind, name).
   The kind is the recognised tag before the last [-]; otherwise the whole
   component is the name (so operator-ish names aren't mis-split). *)
let split_kind comp =
  match String.rindex_opt comp '-' with
  | Some j -> (
      let k = String.sub comp 0 j in
      match List.assoc_opt k kind_table with
      | Some canon ->
          let name = String.sub comp (j + 1) (String.length comp - j - 1) in
          (k ^ "-", Some canon, name)
      | None -> ("", None, comp))
  | None -> ("", None, comp)

(* {1 Enumerating the children of a located item} *)

(* Section headings in a doc comment are referenceable as [label-]/[section-]
   children of their enclosing module or page. Both explicit ([{1:foo …}]) and
   auto-generated labels are offered — both resolve. *)
let labels (docs : Odoc_model.Comment.docs) =
  List.filter_map
    (fun el ->
      match el.Odoc_model.Location_.value with
      | `Heading (_, label, _) -> Some (nm (label :> Id.t), "label")
      | _ -> None)
    docs.elements

let sig_children (sg : Lang.Signature.t) =
  let rec items acc its = List.fold_left item acc its
  and item acc (it : Lang.Signature.item) =
    match it with
    | Module (_, m) -> (nm (m.id :> Id.t), "module") :: acc
    | ModuleType mt -> (nm (mt.id :> Id.t), "module-type") :: acc
    | ModuleSubstitution ms -> (nm (ms.id :> Id.t), "module") :: acc
    | ModuleTypeSubstitution mts -> (nm (mts.id :> Id.t), "module-type") :: acc
    | Type (_, t) -> (nm (t.id :> Id.t), "type") :: acc
    | TypeSubstitution t -> (nm (t.id :> Id.t), "type") :: acc
    | TypExt te ->
        List.fold_left
          (fun acc (c : Lang.Extension.Constructor.t) ->
            (nm (c.id :> Id.t), "extension") :: acc)
          acc te.constructors
    | Exception e -> (nm (e.id :> Id.t), "exception") :: acc
    | Value v -> (nm (v.id :> Id.t), "val") :: acc
    | Class (_, c) -> (nm (c.id :> Id.t), "class") :: acc
    | ClassType (_, ct) -> (nm (ct.id :> Id.t), "class-type") :: acc
    | Include i -> items acc i.expansion.content.items
    | Comment (`Docs d) -> List.rev_append (labels d) acc
    | Open _ | Comment `Stop -> acc
  in
  (* Members and the section labels in the top-comment and standalone comments. *)
  List.rev_append (labels sg.doc) (List.rev (items [] sg.items))

let class_children (cs : Lang.ClassSignature.t) =
  List.filter_map
    (fun (it : Lang.ClassSignature.item) ->
      match it with
      | Method m -> Some (nm (m.id :> Id.t), "method")
      | InstanceVariable v -> Some (nm (v.id :> Id.t), "instance-variable")
      | Constraint _ | Inherit _ | Comment _ -> None)
    cs.items

let type_children (t : Lang.TypeDecl.t) =
  match t.representation with
  | Some (Variant cs) ->
      List.map (fun (c : Lang.TypeDecl.Constructor.t) ->
          (nm (c.id :> Id.t), "constructor"))
        cs
  | Some (Record fs) ->
      List.map (fun (f : Lang.TypeDecl.Field.t) -> (nm (f.id :> Id.t), "field"))
        fs
  | Some (Record_unboxed_product fs) ->
      List.map
        (fun (f : Lang.TypeDecl.UnboxedField.t) -> (nm (f.id :> Id.t), "field"))
        fs
  | Some Extensible | None -> []

let children_of : Refs.located -> (string * string) list = function
  | Lmodule (_, Some sg) -> sig_children sg
  | Lclass (_, Some cs) -> class_children cs
  | Ltype (_, t) -> type_children t
  | Lmodule (_, None) | Lclass (_, None) | Lleaf _ -> []

(* Filter children by name prefix (and kind, if a tag was given), dedup by name,
   and re-attach the verbatim stem and kind prefix. *)
let format ~stem ~kindpfx ~name_prefix ~kind_filter children =
  children
  |> List.filter (fun (n, k) ->
         starts_with ~prefix:name_prefix n
         && match kind_filter with None -> true | Some kf -> kf = k)
  |> List.map fst
  |> List.sort_uniq compare
  |> List.map (fun n -> stem ^ kindpfx ^ n)

(* The children of the item named by a (complete) context reference. *)
let context_children scan ~context =
  match Refs.resolve_to_id scan context with
  | Error _ -> []
  | Ok id -> (
      match Refs.owning_content scan ~ref_str:context id with
      | Ok (Unit_content cu) -> (
          match Refs.find id cu with
          | Some loc -> children_of loc
          | None -> [])
      | Ok (Page_content page) -> labels page.Lang.Page.content
      | Ok _ | Error _ -> [])

(* {1 The three forms} *)

let complete_dotted scan ~stem ~context ~comp =
  let kindpfx, kind_filter, name_prefix = split_kind comp in
  format ~stem ~kindpfx ~name_prefix ~kind_filter (context_children scan ~context)

let complete_toplevel scan ~comp =
  let kindpfx, kind_filter, name_prefix = split_kind comp in
  (* Every compilation unit (a module), minus the page/impl/asset files. *)
  let units =
    Hashtbl.fold (fun k _ acc -> k :: acc) scan.Refs.odocl []
    |> List.filter (fun k ->
           (not (hidden k))
           && not
                (List.exists
                   (fun p -> starts_with ~prefix:p k)
                   [ "Page-"; "Impl-"; "Asset-" ]))
    |> List.map (fun k -> (k, "module"))
  in
  (* Members of the open modules, so bare [List] / [print_endline] complete. *)
  let opened =
    List.concat_map
      (fun m -> context_children scan ~context:m)
      Refs.default_open
  in
  format ~stem:"" ~kindpfx ~name_prefix ~kind_filter (units @ opened)

(* A path segment's name from a directory entry: a subdir is a sub-package /
   sub-library / module; a [foo.odoc] is unit [Foo]; a [page-x.odoc] is page
   [x]; implementations and assets are not referenceable here. *)
let entry_name p ~is_dir =
  let base = Fpath.basename p in
  if is_dir then Some base
  else if Filename.extension base <> ".odoc" then None
  else
    let stem = Filename.remove_extension base in
    if starts_with ~prefix:"page-" stem then
      Some (String.sub stem 5 (String.length stem - 5))
    else if starts_with ~prefix:"impl-" stem || starts_with ~prefix:"asset-" stem
    then None
    else if stem = "" then None
    else Some (String.capitalize_ascii stem)

let complete_root_names scan ~stem ~comp =
  List.map fst scan.Refs.page_roots @ List.map fst scan.Refs.lib_roots
  |> List.filter (starts_with ~prefix:comp)
  |> List.sort_uniq compare
  |> List.map (fun n -> stem ^ n)

(* List a directory's entries (units, pages, sub-dirs) as path segments. *)
let complete_dir dir ~stem ~comp =
  match Bos.OS.Dir.contents dir with
  | Error _ -> []
  | Ok entries ->
      List.filter_map
        (fun p ->
          let is_dir =
            match Bos.OS.Dir.exists p with Ok b -> b | _ -> false
          in
          entry_name p ~is_dir)
        entries
      |> List.filter (fun n -> (not (hidden n)) && starts_with ~prefix:comp n)
      |> List.sort_uniq compare
      |> List.map (fun n -> stem ^ n)

let complete_path scan ~stem ~comp =
  (* [stem] is the path up to and including the last [/]. The first segment is a
     root name ([odoc], [odoc.model], …); any further segments are literal
     sub-directories beneath it ([/odoc/deprecated/] -> odoc/odoc/deprecated). *)
  match String.split_on_char '/' stem with
  | "" :: "" :: _ -> [] (* "//…" current package: out of scope *)
  | "" :: rest -> (
      match List.filter (( <> ) "") rest with
      | [] -> complete_root_names scan ~stem ~comp (* "/" *)
      | root :: subs -> (
          match
            List.assoc_opt root (scan.Refs.page_roots @ scan.Refs.lib_roots)
          with
          | None -> []
          | Some base ->
              if List.exists (fun s -> s = "." || s = "..") subs then []
              else
                let dir = List.fold_left (fun d s -> Fpath.(d / s)) base subs in
                complete_dir dir ~stem ~comp))
  | _ -> [] (* "./…" relative, or no leading slash: out of scope *)

(* {1 Splitting and dispatch} *)

(* Index of the rightmost top-level [.] or [/] (outside parens/quotes, scanned
   right-to-left as odoc's tokenizer does), or [None]. *)
let last_sep s =
  let rec loop i depth in_quote =
    if i < 0 then None
    else
      let c = s.[i] in
      if in_quote then loop (i - 1) depth (c <> '"')
      else
        match c with
        | '"' -> loop (i - 1) depth true
        | ')' -> loop (i - 1) (depth + 1) false
        | '(' -> loop (i - 1) (max 0 (depth - 1)) false
        | ('.' | '/') when depth = 0 -> Some (i, c)
        | _ -> loop (i - 1) depth false
  in
  loop (String.length s - 1) 0 false

let complete sw input =
  let scan = Refs.scan sw in
  let n = String.length input in
  let after i = String.sub input (i + 1) (n - i - 1) in
  match last_sep input with
  | None -> complete_toplevel scan ~comp:input
  | Some (i, '.') ->
      complete_dotted scan
        ~stem:(String.sub input 0 (i + 1))
        ~context:(String.sub input 0 i) ~comp:(after i)
  | Some (i, _ (* '/' *)) ->
      complete_path scan ~stem:(String.sub input 0 (i + 1)) ~comp:(after i)