open! Dune_exe
open! Dune_exe__Rpc__
open Import
module Client = Root.Rpc.Client
module Rpc_error = Dune_rpc.Response.Error

let no_running_server_error (lock_held_by : Global_lock.Lock_held_by.t) =
  match lock_held_by with
  | Global_lock.Lock_held_by.Unknown ->
    User_error.make [ Pp.paragraph "RPC server not running." ]
  | Pid_from_lockfile _ ->
    User_error.make
      [ Pp.paragraph "Another Dune instance is currently running. Aborting..." ]
;;

let active_server ?(lock_held_by = Global_lock.Lock_held_by.Unknown) () =
  match Dune_rpc_impl.Where.get () with
  | Some p -> Ok p
  | None -> Error (no_running_server_error lock_held_by)
;;

let active_server_exn () = active_server () |> User_error.ok_exn

(* cwong: Should we put this into [dune-rpc]? *)
let interpret_kind = function
  | Rpc_error.Invalid_request -> "Invalid_request"
  | Code_error -> "Code_error"
  | Connection_dead -> "Connection_dead"
;;

let raise_rpc_error (e : Rpc_error.t) =
  User_error.raise
    [ Pp.paragraph "Server returned error: "
    ; Pp.paragraphf "%s (error kind: %s)" e.message (interpret_kind e.kind)
      (* CR-soon ElectreAAS: Should we be printing the payload? *)
    ]
;;

let request_exn client request arg =
  let open Fiber.O in
  let* decl =
    Client.Versioned.prepare_request client (Dune_rpc.Decl.Request.witness request)
  in
  match decl with
  | Ok decl ->
    Client.request client decl arg
    >>| (function
     | Ok response -> response
     | Error e -> raise_rpc_error e)
  | Error e -> raise (Dune_rpc.Version_error.E e)
;;

let notify_exn client notification arg =
  let open Fiber.O in
  let* res =
    Client.Versioned.prepare_notification
      client
      (Dune_rpc.Decl.Notification.witness notification)
  in
  match res with
  | Ok decl -> Client.notification client decl arg
  | Error e -> raise (Dune_rpc.Version_error.E e)
;;

let client_term builder f =
  let builder = Common.Builder.forbid_builds builder in
  let builder = Common.Builder.disable_log_file builder in
  let common, config = Common.init builder in
  Scheduler_setup.go_with_rpc_server ~common ~config f
;;

let wait_term =
  let doc = "Poll until server starts listening and then establish connection." in
  Arg.(value & flag & info [ "wait" ] ~doc:(Some doc))
;;

let establish_connection ?(lock_held_by = Global_lock.Lock_held_by.Unknown) () =
  match active_server ~lock_held_by () with
  | Error e -> Fiber.return (Error e)
  | Ok where -> Client.Connection.connect where
;;

let establish_connection_with_retry () =
  let open Fiber.O in
  let pause_between_retries_s = Time.Span.of_secs 0.2 in
  let rec loop () =
    establish_connection ()
    >>= function
    | Ok x -> Fiber.return x
    | Error _ ->
      let* () = Scheduler.sleep pause_between_retries_s in
      loop ()
  in
  loop ()
;;

let establish_client_session ~wait ~lock_held_by =
  let open Fiber.O in
  if wait
  then establish_connection_with_retry ()
  else establish_connection ~lock_held_by () >>| User_error.ok_exn
;;

let prepare_targets targets =
  List.map targets ~f:(fun target ->
    let sexp = Dune_lang.Dep_conf.encode target in
    Dune_lang.to_string sexp)
;;

let warn_ignore_arguments lock_held_by =
  User_warning.emit
    [ Pp.paragraphf
        "Your build request is being forwarded to a running Dune instance%s. Note that \
         certain command line arguments may be ignored."
        (match lock_held_by with
         | Global_lock.Lock_held_by.Unknown -> ""
         | Pid_from_lockfile pid -> sprintf " (pid: %d)" pid)
    ]
;;

let should_warn ~warn_forwarding builder =
  (not Execution_env.inside_dune)
  && warn_forwarding
  && not (Common.Builder.equal builder Common.Builder.default)
;;

let send_request ~f connection name =
  Dune_rpc_impl.Client.client
    connection
    (Dune_rpc.Initialize.Request.create ~id:(Dune_rpc.Id.make (Sexp.Atom name)))
    ~f
;;

let fire_request
      ~name
      ~wait
      ?(warn_forwarding = true)
      ?(lock_held_by = Global_lock.Lock_held_by.Unknown)
      builder
      request
      arg
  =
  let open Fiber.O in
  let* connection = establish_client_session ~wait ~lock_held_by in
  if should_warn ~warn_forwarding builder then warn_ignore_arguments lock_held_by;
  send_request connection name ~f:(fun client -> request_exn client request arg)
;;

let fire_notification
      ~name
      ~wait
      ?(warn_forwarding = true)
      ?(lock_held_by = Global_lock.Lock_held_by.Unknown)
      builder
      notification
      arg
  =
  let open Fiber.O in
  let* connection = establish_client_session ~wait ~lock_held_by in
  if should_warn ~warn_forwarding builder then warn_ignore_arguments lock_held_by;
  send_request connection name ~f:(fun client -> notify_exn client notification arg)
;;

let print_err_warn (num_errors, num_warnings) =
  let enumeration =
    let report_one what count =
      if count = 0
      then []
      else (
        let plural = if count = 1 then "" else "s" in
        [ sprintf "%d %s%s" count what plural ])
    in
    [ report_one "error" num_errors; report_one "warning" num_warnings ]
    |> List.concat
    |> String.enumerate_and
  in
  if num_errors >= 1
  then User_error.raise [ Pp.textf "Build failed with %s." enumeration ]
  else User_warning.emit [ Pp.textf "Build completed with %s." enumeration ]
;;

let wrap_build_outcome_exn ~print_on_success build_outcome =
  match build_outcome with
  | Dune_rpc.Build_outcome_with_diagnostics.Success ->
    if print_on_success
    then Console.print [ Pp.text "Success" |> Pp.tag User_message.Style.Success ]
  | Failure errors ->
    let counts =
      List.fold_left
        errors
        ~init:(0, 0)
        ~f:(fun (errors, warnings) { Dune_rpc.Compound_user_error.main; severity; _ } ->
          match severity with
          | Error ->
            Console.print_user_message main;
            errors + 1, warnings
          | Warning ->
            User_warning.emit_message main;
            errors, warnings + 1)
    in
    (match counts with
     | 0, 0 ->
       Code_error.raise
         "Build via RPC failed, but the RPC server did not send an error message."
         []
     | _ -> print_err_warn counts)
;;
