open! Dune_pkg__
open Import
open Fiber.O

module Curl = struct
  let bin =
    lazy
      (match Bin.which ~path:(Env_path.path Env.initial) "curl" with
       | Some p -> p
       | None ->
         let curl = User_message.command "curl" in
         let sep = Pp.space in
         User_error.raise
           ~hints:
             [ Pp.concat
                 ~sep
                 [ Pp.text "Install"; curl; Pp.text "with your system package manager." ]
             ]
           [ Pp.concat
               ~sep
               [ Pp.text "The program"
               ; curl
               ; Pp.text "does not appear to be installed. Dune uses"
               ; curl
               ; Pp.text "to download packages. Dune requires that the"
               ; curl
               ; Pp.text
                   "executable be located in one of the directories listed in the PATH \
                    variable."
               ]
           ])
  ;;

  let user_agent =
    lazy
      (let base = "dune" in
       match Build_info.V1.version () with
       | None -> base
       | Some v -> base ^ "." ^ Build_info.V1.Version.to_string v)
  ;;

  let curl_features_regex =
    (* If these features are present, then --compressed is supported *)
    let features = [ "libz"; "brotli"; "zstd" ] in
    Re.compile
      Re.(seq [ bol; str "Features:"; rep (first (alt (List.map ~f:str features))) ])
  ;;

  let compressed_supported =
    (* We check if curl supports --compressed by running curl -V and checking if
       the output contains the features we need. *)
    Fiber.Lazy.create (fun () ->
      let+ lines, _ =
        let stderr_to =
          Process.Io.make_stderr
            ~output_on_success:Swallow
            ~output_limit:Dune_engine.Execution_parameters.Action_output_limit.default
        in
        Process.run_capture_lines
          Return
          ~stderr_to
          ~display:Quiet
          (Lazy.force bin)
          [ "-V" ]
      in
      match List.find_map lines ~f:(Re.exec_opt curl_features_regex) with
      | Some group -> Re.Group.test group 0
      | None -> false)
  ;;

  let run ~url ~temp_dir ~output =
    let* compressed_supported = Fiber.Lazy.force compressed_supported in
    let args =
      List.flatten
        [ [ "-L"
          ; "-s"
          ; "--user-agent"
          ; Lazy.force user_agent
          ; "--write-out"
          ; "\"%{http_code}\"" (* This arg must be quoted to work on windows *)
          ; "-o"
          ; Path.to_string output
          ]
        ; (if compressed_supported then [ "--compressed" ] else [])
        ; [ "--"; url ]
        ]
    in
    let stderr = Path.relative temp_dir "curl.stderr" in
    let+ http_code, exit_code =
      let stderr_to = Process.Io.file stderr Out in
      Process.run_capture_line
        Return
        ~stderr_to
        ~display:!Dune_engine.Clflags.display
        (Lazy.force bin)
        args
    in
    if exit_code <> 0
    then (
      let stderr =
        match Io.read_file stderr with
        | s ->
          Fpath.unlink_no_err (Path.to_string stderr);
          [ Pp.text s ]
        | exception s ->
          [ Pp.textf
              "Failed to read stderr from file %s"
              (Path.to_string_maybe_quoted stderr)
          ; Exn.pp s
          ]
      in
      Error
        (User_message.make
           ([ Pp.concat
                ~sep:Pp.space
                [ User_message.command "curl"
                ; Pp.textf "returned an invalid error code %d" exit_code
                ]
            ]
            @ stderr)))
    else (
      Fpath.unlink_no_err (Path.to_string stderr);
      match
        let open Option.O in
        let suffix = {|"|} in
        let prefix = suffix in
        String.drop_prefix_and_suffix ~prefix ~suffix http_code >>= Int.of_string
      with
      | None ->
        Error
          (User_message.make
             [ Pp.concat
                 ~sep:Pp.space
                 [ User_message.command "curl"
                 ; Pp.textf "returned an HTTP code we don't understand: %S" http_code
                 ]
             ])
      | Some http_code ->
        if http_code = 200
        then Ok ()
        else
          Error (User_message.make [ Pp.textf "Download failed with code %d" http_code ]))
  ;;
end

type failure =
  | Checksum_mismatch of Checksum.t
  | Unavailable of User_message.t option

let unpack_archive ~archive_driver ~target ~archive =
  Archive_driver.extract archive_driver ~archive ~target
  >>| Result.map_error ~f:(fun () ->
    Pp.textf "Unable to extract %s" (Path.to_string_maybe_quoted archive))
;;

let check_checksum checksum path =
  let checksum_error =
    match checksum with
    | None -> None
    | Some expected ->
      OpamHash.mismatch (Path.to_string path) (Checksum.to_opam_hash expected)
  in
  match checksum_error with
  | Some c -> Error (Checksum_mismatch (Checksum.of_opam_hash c))
  | None -> Ok ()
;;

let with_download url checksum ~target ~f =
  let url = OpamUrl.to_string url in
  let temp_dir =
    let prefix = "dune" in
    let suffix = Filename.basename url in
    Temp_dir.dir_for_target ~target ~prefix ~suffix
  in
  let output = Path.relative temp_dir "download" in
  Fiber.finalize ~finally:(fun () ->
    Temp.destroy Dir temp_dir;
    Fiber.return ())
  @@ fun () ->
  Curl.run ~temp_dir ~url ~output
  >>= function
  | Error message -> Fiber.return @@ Error (Unavailable (Some message))
  | Ok () ->
    (match check_checksum checksum output with
     | Ok () -> f output
     | Error _ as e -> Fiber.return e)
;;

let fetch_curl ~unpack:unpack_flag ~checksum ~target (url : OpamUrl.t) =
  with_download url checksum ~target ~f:(fun output ->
    match unpack_flag with
    | false ->
      Unix.rename (Path.to_string output) (Path.to_string target);
      Fiber.return @@ Ok ()
    | true ->
      unpack_archive
        ~archive_driver:
          (Archive_driver.choose_for_filename_default_to_tar (OpamUrl0.to_string url))
        ~target
        ~archive:output
      >>| (function
       | Ok () -> Ok ()
       | Error msg ->
         let exn =
           User_message.make
             [ Pp.textf
                 "Failed to unpack archive downloaded from %s"
                 (OpamUrl.to_string url)
             ; Pp.text "Reason:"
             ; msg
             ]
         in
         Error (Unavailable (Some exn))))
;;

let fetch_git rev_store ~target ~url:(url_loc, url) =
  OpamUrl.resolve url ~loc:url_loc rev_store
  >>= (function
   | Error _ as e -> Fiber.return e
   | Ok r -> OpamUrl.fetch_revision url ~loc:url_loc r rev_store)
  >>= function
  | Error msg -> Fiber.return @@ Error (Unavailable (Some msg))
  | Ok at_rev ->
    let+ res = Rev_store.At_rev.check_out at_rev ~target in
    Ok res
;;

let fetch_local ~checksum ~target (url, url_loc) =
  if not (OpamUrl.is_local url)
  then Code_error.raise "fetch_local: url should be file://" [ "url", OpamUrl.to_dyn url ];
  let path =
    match OpamUrl.classify url url_loc with
    | `Path p -> p
    | `Git | `Archive ->
      Code_error.raise "fetch_local: not a path" [ "url", OpamUrl.to_dyn url ]
  in
  match check_checksum checksum path with
  | Error _ as e -> Fiber.return e
  | Ok () ->
    let+ unpack_result =
      unpack_archive
        ~archive_driver:
          (Archive_driver.choose_for_filename_default_to_tar (OpamUrl0.to_string url))
        ~target
        ~archive:path
    in
    Result.map_error unpack_result ~f:(fun pp ->
      Unavailable (Some (User_message.make [ Pp.text "Could not unpack:"; pp ])))
;;

let fetch ~unpack ~checksum ~target ~url:(url_loc, url) =
  let event =
    Dune_trace.(
      Out.start (global ()) (fun () ->
        Dune_trace.Event.Async.fetch
          ~url:(OpamUrl.to_string url)
          ~target
          ~checksum:(Option.map ~f:Checksum.to_string checksum)))
  in
  let unsupported_backend s =
    User_error.raise ~loc:url_loc [ Pp.textf "Unsupported backend: %s" s ]
  in
  Fiber.finalize
    ~finally:(fun () ->
      Option.iter (Dune_trace.global ()) ~f:(fun trace ->
        Dune_trace.Out.finish trace event);
      Fiber.return ())
    (fun () ->
       match url.backend with
       | `git ->
         let* rev_store = Rev_store.get in
         fetch_git rev_store ~target ~url:(url_loc, url)
       | `http -> fetch_curl ~unpack ~checksum ~target url
       | `rsync ->
         if not unpack
         then
           Code_error.raise "fetch_local: unpack is not set" [ "url", OpamUrl.to_dyn url ];
         fetch_local ~checksum ~target (url, url_loc)
       | `hg -> unsupported_backend "mercurial"
       | `darcs -> unsupported_backend "darcs")
;;

let fetch_without_checksum ~unpack ~target ~url =
  fetch ~unpack ~checksum:None ~url ~target
  >>| function
  | Ok () -> Ok ()
  | Error (Checksum_mismatch _) -> assert false
  | Error (Unavailable message) -> Error message
;;
