open Import

module V1 = struct
  module Path = Path
  module Glob = Dune_glob.V1
  open Protocol

  module Execution_error = struct
    exception E of string

    let raise string = raise (E string)

    let raise_on_fs_error = function
      | Error message -> raise message
      | Ok result -> result
    ;;
  end

  module Fs : sig
    val read_directory : string -> (string list, string) result
    val read_file : string -> (string, string) result
    val write_file : string -> string -> (unit, string) result
  end = struct
    let catch_system_exceptions f ~name =
      try Ok (f ()) with
      | Unix.Unix_error (error, syscall, arg) ->
        let error = Unix_error.Detailed.create error ~syscall ~arg in
        Error (name ^ ": " ^ Unix_error.Detailed.to_string_hum error)
      | Sys_error error -> Error (name ^ ": " ^ error)
    ;;

    let read_directory =
      let rec loop dh acc =
        match Unix.readdir dh with
        | "." | ".." -> loop dh acc
        | s -> loop dh (s :: acc)
        | exception End_of_file -> acc
      in
      fun path ->
        catch_system_exceptions ~name:"read_directory" (fun () ->
          let dh = Unix.opendir path in
          Exn.protect
            ~f:(fun () -> loop dh [] |> List.sort ~compare:String.compare)
            ~finally:(fun () -> Unix.closedir dh))
    ;;

    let read_file path =
      catch_system_exceptions ~name:"read_file" (fun () -> Io.String_path.read_file path)
    ;;

    let write_file path data =
      catch_system_exceptions ~name:"write_file" (fun () ->
        Io.String_path.write_file path data)
    ;;
  end

  module Stage = struct
    type 'a t =
      { action : unit -> 'a
      ; dependencies : Dependency.Set.t
      ; targets : String.Set.t
      }

    let map (t : 'a t) ~f = { t with action = (fun () -> f (t.action ())) }

    let both (t1 : 'a t) (t2 : 'b t) =
      { action = (fun () -> t1.action (), t2.action ())
      ; dependencies = Dependency.Set.union t1.dependencies t2.dependencies
      ; targets = String.Set.union t1.targets t2.targets
      }
    ;;
  end

  (* Construction inspired by free monad. *)
  type 'a t =
    | Pure of 'a
    | Stage of 'a t Stage.t

  let lift_stage stage = Stage (Stage.map stage ~f:(fun a -> Pure a))

  let rec map (t : 'a t) ~f =
    match t with
    | Pure a -> Pure (f a)
    | Stage at -> Stage (Stage.map ~f:(map ~f) at)
  ;;

  let rec stage (t : 'a t) ~f =
    match t with
    | Pure a -> f a
    | Stage at -> Stage (Stage.map ~f:(stage ~f) at)
  ;;

  let return a = Pure a

  let rec both (t1 : 'a t) (t2 : 'b t) =
    match t1, t2 with
    | Pure a1, _ -> map ~f:(fun a2 -> a1, a2) t2
    | _, Pure a2 -> map ~f:(fun a1 -> a1, a2) t1
    | Stage at1, Stage at2 ->
      Stage (Stage.both at1 at2 |> Stage.map ~f:(fun (am1, am2) -> both am1 am2))
  ;;

  let read_file ~path =
    let path = Path.to_string path in
    let action () = Fs.read_file path |> Execution_error.raise_on_fs_error in
    lift_stage
      { action
      ; dependencies = Dependency.Set.singleton (File path)
      ; targets = String.Set.empty
      }
  ;;

  let write_file ~path ~data =
    let path = Path.to_string path in
    let action () = Fs.write_file path data |> Execution_error.raise_on_fs_error in
    lift_stage
      { action; dependencies = Dependency.Set.empty; targets = String.Set.singleton path }
  ;;

  (* TODO jstaron: If program tries to read empty directory, dune does not copy
     it to `_build` so we get a "No such file or directory" error. *)
  let read_directory_with_glob ~path ~glob =
    let path = Path.to_string path in
    let action () =
      Fs.read_directory path
      |> Execution_error.raise_on_fs_error
      |> List.filter ~f:(Glob.test glob)
    in
    lift_stage
      { action
      ; dependencies =
          Dependency.Set.singleton (Glob { path; glob = Glob.to_string glob })
      ; targets = String.Set.empty
      }
  ;;

  let rec run_by_dune t context =
    match t with
    | Pure () -> Context.respond context Done
    | Stage at ->
      let allowed_targets = Context.targets context in
      let disallowed_targets = String.Set.diff at.targets allowed_targets in
      (match String.Set.to_list disallowed_targets with
       | [] -> ()
       | [ t ] ->
         Execution_error.raise
           (Printf.sprintf
              "%s is written despite not being declared as a target in dune file. To \
               fix, add it to target list in dune file."
              t)
       | ts ->
         Execution_error.raise
           (Printf.sprintf
              "Following files were written despite not being declared as targets in \
               dune file:\n\
               %sTo fix, add them to target list in dune file."
              (ts |> String.concat ~sep:"\n")));
      let prepared_dependencies = Context.prepared_dependencies context in
      let required_dependencies =
        Dependency.Set.diff at.dependencies prepared_dependencies
      in
      if Dependency.Set.is_empty required_dependencies
      then run_by_dune (at.action ()) context
      else Context.respond context (Need_more_deps required_dependencies)
  ;;

  (* If executable is not run by dune, assume that all dependencies are already
     prepared and no target checking is done. *)
  let rec run_outside_of_dune t =
    match t with
    | Pure () -> ()
    | Stage at -> run_outside_of_dune (at.action ())
  ;;

  let do_run t =
    match Protocol.Context.create () with
    | Run_outside_of_dune -> run_outside_of_dune t
    | Error message ->
      Execution_error.raise
        (Printf.sprintf
           "Error during communication with dune. %s Did you use different dune version \
            to compile the executable?"
           message)
    | Ok context -> run_by_dune t context
  ;;

  let run t =
    try
      do_run t;
      exit 0
    with
    | Execution_error.E message ->
      prerr_endline message;
      exit 1
  ;;

  module O = struct
    let ( let+ ) at f = map at ~f
    let ( and+ ) = both
  end

  module Private = struct
    module Protocol = Protocol

    let do_run = do_run

    module Execution_error = Execution_error
  end
end

module Private = V1.Private
