open Import
open Memo.O

type t =
  { loc : Loc.t
  ; files : Predicate_lang.Glob.t
  ; libraries : Lib_dep.t list
  ; preprocess : Preprocess.Without_instrumentation.t Preprocess.Per_module.t
  ; preprocessor_deps : Dep_conf.t list
  ; runtime_deps : Dep_conf.t list
  ; cinaps_version : Syntax.Version.t
  ; alias : Alias.Name.t option
  ; link_flags : Dune_lang.Link_flags.Spec.t
  }

let name = "cinaps"
let cinaps_alias = Alias.Name.of_string name

include Stanza.Make (struct
    type nonrec t = t

    include Poly
  end)

let syntax =
  Dune_lang.Syntax.create
    ~name:(Syntax.Name.parse name)
    ~desc:"the cinaps extension"
    [ (1, 0), `Since (1, 11)
    ; (1, 1), `Since (3, 5)
    ; (1, 2), `Since (3, 7)
    ; (1, 3), `Since (3, 8)
    ]
;;

let decode =
  let open Dune_lang.Decoder in
  fields
    (let+ loc = loc
     and+ files = field "files" Predicate_lang.Glob.decode ~default:Predicate_lang.true_
     and+ preprocess, preprocessor_deps = Preprocess.preprocess_fields
     and+ libraries =
       field "libraries" (Lib_dep.L.decode ~allow_re_export:false) ~default:[]
     and+ runtime_deps =
       field
         ~default:[]
         "runtime_deps"
         (Dune_lang.Syntax.since syntax (1, 1) >>> repeat Dep_conf.decode)
     and+ cinaps_version = Dune_lang.Syntax.get_exn syntax
     and+ alias = field_o "alias" Dune_lang.Alias.decode
     and+ link_flags =
       Dune_lang.Link_flags.Spec.decode
         ~check:(Some (Dune_lang.Syntax.since syntax (1, 3)))
     (* TODO use this field? *)
     and+ _flags = Dune_lang.Ocaml_flags.Spec.decode in
     { loc
     ; files
     ; libraries
     ; preprocess
     ; preprocessor_deps
     ; runtime_deps
     ; cinaps_version
     ; alias
     ; link_flags
     })
;;

let () =
  let open Dune_lang.Decoder in
  Dune_project.Extension.register_simple syntax (return [ name, decode_stanza decode ])
;;

let flags = Ocaml_flags.of_list [ "-w"; "-24" ]
let for_ = Compilation_mode.Ocaml

let gen_rules sctx t ~dir ~scope =
  let digest_input_repr =
    Repr.variant
      "cinaps-digest-input"
      [ Repr.case
          "No_files"
          Repr.(
            triple
              (list Lib_dep.repr)
              (Preprocess.Per_module.repr Preprocess.Without_instrumentation.repr)
              (list Dep_conf.repr))
          ~proj:(function
            | `No_files x -> Some x
            | `With_files _ -> None)
      ; Repr.case
          "With_files"
          Repr.(
            pair
              (list Path.Build.repr)
              (triple
                 (list Lib_dep.repr)
                 (Preprocess.Per_module.repr Preprocess.Without_instrumentation.repr)
                 (list Dep_conf.repr)))
          ~proj:(function
            | `With_files x -> Some x
            | `No_files _ -> None)
      ]
  in
  let loc = t.loc in
  (* Files checked by cinaps *)
  let* cinapsed_files =
    Source_tree.files_of (Path.Build.drop_build_context_exn dir)
    >>| Path.Source.Set.to_list
    >>| List.filter_map ~f:(fun p ->
      if
        Predicate_lang.Glob.test
          t.files
          (Path.Source.basename p)
          ~standard:Predicate_lang.true_
      then
        Some
          (Path.Build.append_source (Super_context.context sctx |> Context.build_dir) p)
      else None)
  in
  let cinaps_dir =
    let stamp =
      let digest =
        Digest.repr
          digest_input_repr
          (match cinapsed_files with
           | [] ->
             `No_files
               ( t.libraries
               , t.preprocess
               , List.map t.preprocessor_deps ~f:Dep_conf.remove_locs )
           | _ :: _ ->
             `With_files
               ( cinapsed_files
               , ( t.libraries
                 , t.preprocess
                 , List.map t.preprocessor_deps ~f:Dep_conf.remove_locs ) ))
      in
      String.take (Digest.to_string digest) 8
    in
    Path.Build.relative dir ("." ^ name ^ "." ^ stamp)
  in
  let main_module_name = Module_name.of_checked_string name in
  let module_ =
    Module.generated ~kind:Impl ~for_ [ main_module_name ] ~src_dir:cinaps_dir
  in
  let* () =
    let cinaps_ml =
      Module.source ~ml_kind:Ml_kind.Impl module_
      |> Option.value_exn
      |> Module.File.path
      |> Path.as_in_build_dir_exn
    in
    (* Ask cinaps to produce a .ml file to build *)
    let sandbox =
      if t.cinaps_version >= (1, 1)
      then Sandbox_config.needs_sandboxing
      else Sandbox_config.default
    in
    Super_context.add_rule
      sctx
      ~loc:t.loc
      ~dir
      (let prog =
         Super_context.resolve_program
           sctx
           ~dir
           ~where:Original_path
           ~loc:(Some loc)
           name
           ~hint:"opam install cinaps"
       in
       Command.run_dyn_prog
         ~dir:(Path.build dir)
         prog
         ~sandbox
         [ A "-staged"; Target cinaps_ml; Deps (List.map cinapsed_files ~f:Path.build) ])
  and* expander = Super_context.expander sctx ~dir in
  let compile_info =
    let dune_version = Scope.project scope |> Dune_project.dune_version in
    let names = Nonempty_list.[ t.loc, name ] in
    Lib.DB.resolve_user_written_deps
      (Scope.libs scope)
      (`Exe names)
      (Lib_dep.Direct (loc, Lib_name.of_string "cinaps.runtime") :: t.libraries)
      ~allow_unused_libraries:[]
      ~pps:(Preprocess.Per_module.pps t.preprocess)
      ~dune_version
      ~allow_overlaps:false
      ~forbidden_libraries:[]
  in
  let* (_ : Exe.dep_graphs) =
    let* cctx =
      let* modules =
        let* preprocess =
          Pp_spec_rules.make
            sctx
            ~dir
            ~expander
            ~lint:(Preprocess.Per_module.no_preprocessing ())
            ~preprocess:t.preprocess
            ~preprocessor_deps:t.preprocessor_deps
            ~instrumentation_deps:[]
            ~lib_name:None
            ~scope
        in
        Pp_spec.pp_module preprocess module_ >>| Modules.With_vlib.singleton_exe
      in
      let requires_compile = Lib.Compile.direct_requires compile_info ~for_ in
      let requires_link = Lib.Compile.requires_link compile_info ~for_ in
      let obj_dir = Obj_dir.make_exe ~dir:cinaps_dir ~name in
      Compilation_context.create
        for_
        ~super_context:sctx
        ~scope
        ~obj_dir
        ~modules
        ~opaque:(Explicit false)
        ~requires_compile
        ~requires_link
        ~flags
        ~js_of_ocaml:(Js_of_ocaml.Mode.Pair.make None)
        ~melange_package_name:None
        ~package:None
        ~bin_annot:false
    in
    let link_args =
      let open Action_builder.O in
      let+ link_args =
        Ocaml_flags_db.link_flags sctx ~dir t.link_flags
        |> Action_builder.of_memo
        >>= Link_flags.get ~use_standard_cxx_flags:false
      in
      Command.Args.As link_args
    in
    Exe.build_and_link
      cctx
      ~link_args
      ~program:{ name; main_module_name; loc }
      ~linkages:[ Exe.Linkage.native_or_custom (Compilation_context.ocaml cctx) ]
      ~promote:None
  in
  let cinaps_alias = Alias.make ~dir @@ Option.value t.alias ~default:cinaps_alias in
  let* () =
    let action =
      let open Action_builder.O in
      let cinaps_exe =
        let cinaps_exe = Path.Build.relative cinaps_dir (name ^ ".exe") in
        Path.build cinaps_exe
      in
      let runtime_deps, sandbox =
        let sandbox =
          if t.cinaps_version >= (1, 1)
          then Sandbox_config.needs_sandboxing
          else Sandbox_config.no_special_requirements
        in
        Dep_conf_eval.unnamed sandbox ~expander t.runtime_deps
      in
      let* () = runtime_deps in
      let+ () =
        cinaps_exe :: List.rev_map cinapsed_files ~f:Path.build
        |> Dep.Set.of_files
        |> Action_builder.deps
      in
      Action.Full.make ~sandbox
      @@ Action.chdir
           (Path.build dir)
           (Action.progn
              [ Action.run (Ok cinaps_exe) [ "-diff-cmd"; "-" ]
              ; Action.concurrent
                @@ List.map cinapsed_files ~f:(fun fn ->
                  Action.diff
                    ~optional:true
                    (Path.build fn)
                    (Path.Build.extend_basename fn ~suffix:".cinaps-corrected"))
              ])
    in
    Super_context.add_alias_action sctx ~dir ~loc cinaps_alias action
  in
  match t.alias with
  | Some _ -> Memo.return ()
  | None ->
    Rules.Produce.Alias.add_deps
      (Alias.make Alias0.runtest ~dir)
      (Alias_builder.alias cinaps_alias)
;;
