open Import

(* we need to convince ocamldep that we don't depend on the menhir rules *)
module Menhir = struct end
open Memo.O

let loc_of_dune_file st_dir =
  (match
     let open Option.O in
     let* dune_file = Source_tree.Dir.dune_file st_dir in
     (* TODO not really correct. we need to know the [(subdir ..)] that introduced this *)
     Source.Dune_file.path dune_file
   with
   | Some s -> s
   | None -> Path.Source.relative (Source_tree.Dir.path st_dir) "_unknown_")
  |> Path.source
  |> Loc.in_file
;;

type t =
  { kind : kind
  ; dir : Path.Build.t
  ; text_files : Filename.Set.t
  ; foreign_sources : Foreign_sources.t Memo.Lazy.t
  ; mlds : (Documentation.t * Doc_sources.mld list) list Memo.Lazy.t
  ; coq : Coq_sources.t Memo.Lazy.t
  ; rocq : Rocq_sources.t Memo.Lazy.t
  ; ocaml : Ml_sources.t Memo.Lazy.t
  ; melange : Ml_sources.t Memo.Lazy.t
  ; source_dir : Source_tree.Dir.t option
  }

and kind =
  | Standalone
  | Group_root of t list
  | Group_part

let empty kind ~dir ~source_dir =
  { kind
  ; dir
  ; source_dir
  ; text_files = Filename.Set.empty
  ; ocaml = Memo.Lazy.of_val Ml_sources.empty
  ; melange = Memo.Lazy.of_val Ml_sources.empty
  ; mlds = Memo.Lazy.of_val []
  ; foreign_sources = Memo.Lazy.of_val Foreign_sources.empty
  ; coq = Memo.Lazy.of_val Coq_sources.empty
  ; rocq = Memo.Lazy.of_val Rocq_sources.empty
  }
;;

module Standalone_or_root = struct
  type nonrec standalone_or_root =
    { root : t
    ; subdirs : t Path.Build.Map.t
    ; rules : Rules.t
    }

  type nonrec t = { contents : standalone_or_root Memo.Lazy.t }

  let empty ~dir ~source_dir =
    { contents =
        Memo.Lazy.create (fun () ->
          Memo.return
            { root = empty Standalone ~dir ~source_dir
            ; rules = Rules.empty
            ; subdirs = Path.Build.Map.empty
            })
    }
  ;;

  let root t =
    let+ contents = Memo.Lazy.force t.contents in
    contents.root
  ;;

  let subdirs t =
    let+ contents = Memo.Lazy.force t.contents in
    Path.Build.Map.values contents.subdirs
  ;;

  let rules t =
    let+ contents = Memo.Lazy.force t.contents in
    contents.rules
  ;;
end

type triage =
  | Standalone_or_root of Standalone_or_root.t
  | Group_part of Path.Build.t

let dir t = t.dir
let source_dir t = t.source_dir
let coq t = Memo.Lazy.force t.coq
let rocq t = Memo.Lazy.force t.rocq
let ocaml t = Memo.Lazy.force t.ocaml
let melange t = Memo.Lazy.force t.melange

let ml t ~for_ =
  match for_ with
  | Compilation_mode.Ocaml -> ocaml t
  | Melange -> melange t
;;

let dirs t =
  match t.kind with
  | Standalone -> [ t ]
  | Group_root subs -> t :: subs
  | Group_part ->
    Code_error.raise
      "Dir_contents.dirs called on a group part"
      [ "dir", Path.Build.to_dyn t.dir ]
;;

let text_files t = t.text_files
let foreign_sources t = Memo.Lazy.force t.foreign_sources

let mlds t ~(stanza : Documentation.t) =
  let+ map = Memo.Lazy.force t.mlds in
  match
    List.find_map map ~f:(fun (stanza', x) ->
      Option.some_if (Loc.equal stanza.loc stanza'.loc) x)
  with
  | Some x -> x
  | None ->
    Code_error.raise
      "Dir_contents.mlds"
      [ "doc", Loc.to_dyn_hum stanza.loc
      ; ( "available"
        , Dyn.(list Loc.to_dyn_hum)
            (List.map map ~f:(fun ((d : Documentation.t), _) -> d.loc)) )
      ]
;;

module rec Load : sig
  val get : Super_context.t -> dir:Path.Build.t -> t Memo.t
  val triage : Super_context.t -> dir:Path.Build.t -> triage Memo.t
end = struct
  let select_deps_files ~dialects libraries =
    (* Keep non-OCaml `(select ..)` outputs in the directory contents so
       foreign stubs can depend on generated headers and similar files. *)
    List.filter_map libraries ~f:(fun dep ->
      match (dep : Lib_dep.t) with
      | Re_export _ | Direct _ | Instantiate _ -> None
      | Select s ->
        let ext = Path.Local.basename s.result_fn |> Filename.extension in
        (match Filename.Extension.Or_empty.extension ext with
         | None -> Some (Path.Local.to_string s.result_fn)
         | Some ext ->
           (match Dialect.DB.find_by_extension dialects ext with
            | Some _ -> None
            | None -> Some (Path.Local.to_string s.result_fn))))
  ;;

  (* As a side-effect, setup user rules and copy_files rules. *)
  let load_text_files sctx st_dir stanzas ~dir ~src_dir =
    let from_source = Source_tree.Dir.filenames st_dir in
    match stanzas with
    | [] -> Memo.return from_source
    | _ :: _ ->
      (* Interpret a few stanzas in order to determine the list of files generated
         by the user. *)
      let+ generated_files =
        let* expander = Super_context.expander sctx ~dir in
        Memo.parallel_map stanzas ~f:(fun stanza ->
          match Stanza.repr stanza with
          | Coq_stanza.Coqpp.T { modules; _ } ->
            Coq_sources.mlg_files ~sctx ~dir ~modules
            >>| List.rev_map ~f:(fun mlg_file ->
              Path.Build.set_extension mlg_file ~ext:Filename.Extension.ml
              |> Path.Build.basename)
          | Coq_stanza.Extraction.T s ->
            Memo.return (Coq_stanza.Extraction.ml_target_fnames s)
          | Rocq_stanza.Rocqpp.T { modules; _ } ->
            Rocq_sources.mlg_files ~sctx ~dir ~modules
            >>| List.rev_map ~f:(fun mlg_file ->
              Path.Build.set_extension mlg_file ~ext:Filename.Extension.ml
              |> Path.Build.basename)
          | Rocq_stanza.Extraction.T s ->
            Memo.return (Rocq_stanza.Extraction.target_fnames s)
          | Rule_conf.T rule ->
            Simple_rules.user_rule sctx rule ~dir ~expander
            >>| (function
             | None -> []
             | Some targets ->
               (* CR-someday amokhov: Do not ignore directory targets. *)
               Filename.Set.to_list targets.files)
          | Copy_files.T def ->
            Simple_rules.copy_files sctx def ~src_dir ~dir ~expander
            >>| Path.Set.to_list_map ~f:Path.basename
          | Generate_sites_module_stanza.T def ->
            Generate_sites_module_rules.setup_rules sctx ~dir def >>| List.singleton
          | Library.T { buildable; _ }
          | Executables.T { buildable; _ }
          | Tests.T { exes = { buildable; _ }; _ } ->
            let dialects = Source_tree.Dir.project st_dir |> Dune_project.dialects in
            let select_deps_files = select_deps_files ~dialects buildable.libraries in
            let ctypes_files =
              (* Also manually add files generated by ctypes rules. *)
              Option.map buildable.ctypes ~f:Ctypes_field.generated_ml_and_c_files
              |> Option.value ~default:[]
            in
            Memo.return (select_deps_files @ ctypes_files)
          | _ -> Memo.return [])
        >>| List.concat
        >>| Filename.Set.of_list
      in
      Filename.Set.union generated_files from_source
  ;;

  module Key = struct
    module Super_context = Super_context.As_memo_key

    type t = Super_context.t * Path.Build.t

    let to_dyn (sctx, path) =
      Dyn.Tuple [ Super_context.to_dyn sctx; Path.Build.to_dyn path ]
    ;;

    let equal = Tuple.T2.equal Super_context.equal Path.Build.equal
    let hash = Tuple.T2.hash Super_context.hash Path.Build.hash
  end

  let lookup_vlib sctx ~current_dir ~loc ~dir ~for_ =
    match Path.Build.equal current_dir dir with
    | true ->
      User_error.raise
        ~loc
        [ Pp.text
            "Virtual library and its implementation(s) cannot be defined in the same \
             directory"
        ]
    | false -> Load.get sctx ~dir >>= ml ~for_
  ;;

  let human_readable_description dir =
    Pp.textf
      "Computing directory contents of %s"
      (Path.to_string_maybe_quoted (Path.build dir))
  ;;

  let mlds ~sctx ~dir ~dune_file ~files =
    Memo.lazy_ (fun () ->
      let* expander = Super_context.expander sctx ~dir in
      Doc_sources.build_mlds_map dune_file ~dir ~files expander)
  ;;

  let make_standalone sctx st_dir ~dir (d : Dune_file.t) =
    let human_readable_description () = human_readable_description dir in
    { Standalone_or_root.contents =
        Memo.lazy_ ~human_readable_description (fun () ->
          let include_subdirs = Loc.none, Include_subdirs.No in
          let ctx = Super_context.context sctx in
          let lib_config =
            let+ ocaml = Context.ocaml ctx in
            ocaml.lib_config
          in
          let stanzas = Dune_file.stanzas d in
          let project = Dune_file.project d in
          let+ files, rules =
            Rules.collect (fun () ->
              let src_dir = Dune_file.dir d in
              stanzas >>= load_text_files sctx st_dir ~src_dir ~dir)
          in
          let dirs =
            Memo.lazy_ (fun () ->
              let+ stanzas = stanzas in
              Nonempty_list.
                [ { Source_file_dir.dir
                  ; path_to_root = []
                  ; files
                  ; source_dir = Some st_dir
                  ; stanzas
                  }
                ])
          in
          let ml =
            Memo.lazy_ (fun () ->
              let lookup_vlib = lookup_vlib sctx ~current_dir:dir ~for_:Ocaml in
              let loc = loc_of_dune_file st_dir in
              let libs = Scope.DB.find_by_dir dir >>| Scope.libs in
              let* expander = Super_context.expander sctx ~dir
              and* dirs = Memo.Lazy.force dirs in
              Ml_sources.make
                ~expander
                ~libs
                ~project
                ~lib_config
                ~loc
                ~include_subdirs
                ~lookup_vlib
                dirs)
          in
          let mlds = mlds ~sctx ~dir ~dune_file:d ~files in
          { Standalone_or_root.root =
              { kind = Standalone
              ; source_dir = Some st_dir
              ; dir
              ; text_files = files
              ; ocaml = ml
              ; melange = ml
              ; mlds
              ; foreign_sources =
                  Memo.lazy_ (fun () ->
                    let dune_version = Dune_project.dune_version project in
                    let* stanzas = stanzas
                    and* dirs = Memo.Lazy.force dirs in
                    Foreign_sources.make stanzas ~dir ~dune_version ~dirs)
              ; coq =
                  Memo.lazy_ (fun () ->
                    let+ stanzas = stanzas
                    and+ dirs = Memo.Lazy.force dirs in
                    Coq_sources.of_dir stanzas ~dir ~include_subdirs ~dirs)
              ; rocq =
                  Memo.lazy_ (fun () ->
                    let+ stanzas = stanzas
                    and+ dirs = Memo.Lazy.force dirs in
                    Rocq_sources.of_dir stanzas ~dir ~include_subdirs ~dirs)
              }
          ; rules
          ; subdirs = Path.Build.Map.empty
          })
    }
  ;;

  let make_group_root
        sctx
        ~dir
        { Dir_status.Group_root.qualification; dune_file; source_dir; components }
    =
    let include_subdirs =
      let loc, qualif_mode = qualification in
      loc, Include_subdirs.Include qualif_mode
    in
    let loc = loc_of_dune_file source_dir in
    let+ components = components in
    let contents =
      Memo.lazy_
        ~human_readable_description:(fun () -> human_readable_description dir)
        (fun () ->
           let ctx = Super_context.context sctx in
           let stanzas = Dune_file.stanzas dune_file in
           let project = Dune_file.project dune_file in
           let+ (files, subdirs), rules =
             Rules.collect (fun () ->
               Memo.fork_and_join
                 (fun () ->
                    stanzas
                    >>= load_text_files
                          sctx
                          source_dir
                          ~src_dir:(Dune_file.dir dune_file)
                          ~dir)
                 (fun () ->
                    Memo.parallel_map
                      components
                      ~f:(fun { dir; path_to_group_root; source_dir; stanzas } ->
                        let+ files =
                          load_text_files
                            sctx
                            source_dir
                            stanzas
                            ~src_dir:(Source_tree.Dir.path source_dir)
                            ~dir
                        in
                        { Source_file_dir.dir
                        ; path_to_root = path_to_group_root
                        ; files
                        ; source_dir = Some source_dir
                        ; stanzas
                        })))
           in
           let dirs =
             Memo.lazy_ (fun () ->
               let+ stanzas = stanzas in
               Nonempty_list.(
                 { Source_file_dir.dir
                 ; path_to_root = []
                 ; files
                 ; source_dir = Some source_dir
                 ; stanzas
                 }
                 :: subdirs))
           in
           let lib_config =
             let+ ocaml = Context.ocaml ctx in
             ocaml.lib_config
           in
           let ml =
             Memo.lazy_ (fun () ->
               let lookup_vlib = lookup_vlib sctx ~current_dir:dir ~for_:Ocaml in
               let libs = Scope.DB.find_by_dir dir >>| Scope.libs in
               let* expander = Super_context.expander sctx ~dir
               and* dirs = Memo.Lazy.force dirs in
               Ml_sources.make
                 ~expander
                 ~project
                 ~libs
                 ~lib_config
                 ~loc
                 ~lookup_vlib
                 ~include_subdirs
                 dirs)
           in
           let foreign_sources =
             Memo.lazy_ (fun () ->
               let dune_version = Dune_project.dune_version project in
               let* stanzas = stanzas
               and* dirs = Memo.Lazy.force dirs in
               Foreign_sources.make stanzas ~dir ~dune_version ~dirs)
           in
           let coq =
             Memo.lazy_ (fun () ->
               let+ stanzas = stanzas
               and+ dirs = Memo.Lazy.force dirs in
               Coq_sources.of_dir stanzas ~dir ~dirs ~include_subdirs)
           in
           let rocq =
             Memo.lazy_ (fun () ->
               let+ stanzas = stanzas
               and+ dirs = Memo.Lazy.force dirs in
               Rocq_sources.of_dir stanzas ~dir ~dirs ~include_subdirs)
           in
           let mlds = mlds ~sctx ~dir ~dune_file ~files in
           let subdirs =
             List.map
               subdirs
               ~f:
                 (fun
                   { Source_file_dir.dir
                   ; path_to_root = _
                   ; files
                   ; source_dir
                   ; stanzas = _
                   }
                 ->
                 { kind = Group_part
                 ; source_dir
                 ; dir
                 ; text_files = files
                 ; ocaml = ml
                 ; melange = ml
                 ; foreign_sources
                 ; mlds
                 ; coq
                 ; rocq
                 })
           in
           let root =
             { kind = Group_root subdirs
             ; source_dir = Some source_dir
             ; dir
             ; text_files = files
             ; ocaml = ml
             ; melange = ml
             ; foreign_sources
             ; mlds
             ; coq
             ; rocq
             }
           in
           { Standalone_or_root.root
           ; rules
           ; subdirs = Path.Build.Map.of_list_map_exn subdirs ~f:(fun x -> x.dir, x)
           })
    in
    { Standalone_or_root.contents }
  ;;

  let get0_impl (sctx, dir) : triage Memo.t =
    Dir_status.DB.get ~dir
    >>= function
    | Is_component_of_a_group_but_not_the_root { group_root; stanzas = _ } ->
      Memo.return @@ Group_part group_root
    | Generated ->
      Memo.return @@ Standalone_or_root (Standalone_or_root.empty ~dir ~source_dir:None)
    | Lock_dir source_dir | Source_only source_dir ->
      Memo.return
      @@ Standalone_or_root (Standalone_or_root.empty ~dir ~source_dir:(Some source_dir))
    | Standalone (st_dir, d) ->
      Memo.return @@ Standalone_or_root (make_standalone sctx st_dir ~dir d)
    | Group_root root ->
      let+ group_root = make_group_root sctx root ~dir in
      Standalone_or_root group_root
  ;;

  let memo0 =
    Memo.create
      "dir-contents-get0"
      get0_impl
      ~input:(module Key)
      ~human_readable_description:(fun (_, dir) ->
        Pp.textf
          "Computing directory contents of %s"
          (Path.to_string_maybe_quoted (Path.build dir)))
  ;;

  let get sctx ~dir =
    Memo.exec memo0 (sctx, dir)
    >>= function
    | Standalone_or_root { contents } ->
      let+ { root; rules = _; subdirs = _ } = Memo.Lazy.force contents in
      root
    | Group_part group_root ->
      Memo.exec memo0 (sctx, group_root)
      >>= (function
       | Group_part _ -> assert false
       | Standalone_or_root { contents } ->
         let+ { root; rules = _; subdirs = _ } = Memo.Lazy.force contents in
         root)
  ;;

  let triage sctx ~dir = Memo.exec memo0 (sctx, dir)
end

include Load

let modules_of_local_lib sctx lib ~for_ =
  let info = Lib.Local.info lib in
  let dir = Lib_info.src_dir info in
  let* t = get sctx ~dir
  and* libs = Scope.DB.find_by_dir dir >>| Scope.libs in
  ml t ~for_
  >>= Ml_sources.modules
        ~libs
        ~for_:(Library (Lib_info.lib_id info |> Lib_id.to_local_exn))
;;

let modules_of_lib sctx lib ~for_ =
  match
    let info = Lib.info lib in
    Lib_info.modules info ~for_
  with
  | External modules -> Memo.return modules
  | Local ->
    let+ modules = modules_of_local_lib sctx (Lib.Local.of_lib_exn lib) ~for_ in
    Some (Modules.With_vlib.modules modules)
;;

let () =
  Fdecl.set Expander.lookup_artifacts (fun ~dir ->
    let* t =
      Context.DB.by_dir dir >>| Context.name >>= Super_context.find_exn >>= Load.get ~dir
    in
    Memo.Lazy.force t.ocaml >>= Ml_sources.artifacts)
;;
