defmodule Ex.S3.Auth do
  def put_aws_sigv4(request) do
    if s3_options = request.options[:s3] do
      s3_options =
        s3_options
        |> Keyword.put_new(:region, "us-east-1")
        |> Keyword.put_new(:datetime, DateTime.utc_now())
        # aws_credentials returns this key so let's ignore it
        |> Keyword.drop([:credential_provider, :bucket, :endpoint])
        |> Keyword.put(:service, :s3)

      Req.Request.validate_options(s3_options, [
        :access_key_id,
        :secret_access_key,
        :token,
        :service,
        :region,
        :datetime,

        # for req_s3
        :expires
      ])

      unless s3_options[:access_key_id] do
        raise ArgumentError, "missing :access_key_id in :s3 option"
      end

      unless s3_options[:secret_access_key] do
        raise ArgumentError, "missing :secret_access_key in :s3 option"
      end

      {body, options} =
        case request.body do
          nil ->
            {"", []}

          iodata when is_binary(iodata) or is_list(iodata) ->
            {iodata, []}

          _enumerable ->
            if Req.Request.get_header(request, "content-length") == [] do
              raise "content-length header must be explicitly set when streaming request body"
            end

            {"", [body_digest: "UNSIGNED-PAYLOAD"]}
        end

      request = Req.Request.put_new_header(request, "host", request.url.host)

      headers = for {name, values} <- request.headers, value <- values, do: {name, value}

      headers =
        aws_sigv4_headers(
          s3_options ++
            [
              method: request.method,
              url: to_string(request.url),
              headers: headers,
              body: body
            ] ++ options
        )

      request
      |> Req.merge(headers: headers)
      |> Req.Request.append_response_steps(s3_decode: &decode_body/1)
    else
      request
    end
  end

  defp decode_body({request, response}) do
    if request.method in [:get, :head] and
         request.options[:decode_body] != false and
         request.options[:raw] != true and
         match?(["application/xml" <> _], response.headers["content-type"]) do
      response = update_in(response.body, &Ex.S3.XML.parse_s3/1)
      {request, response}
    else
      {request, response}
    end
  end

  @doc """
  Create AWS Signature v4.

  https://docs.aws.amazon.com/AmazonS3/latest/API/sig-v4-header-based-auth.html
  """
  def aws_sigv4_headers(options) do
    {access_key_id, options} = Keyword.pop!(options, :access_key_id)
    {secret_access_key, options} = Keyword.pop!(options, :secret_access_key)
    {security_token, options} = Keyword.pop(options, :token)
    {region, options} = Keyword.pop!(options, :region)
    {service, options} = Keyword.pop!(options, :service)
    {datetime, options} = Keyword.pop!(options, :datetime)
    {method, options} = Keyword.pop!(options, :method)
    {url, options} = Keyword.pop!(options, :url)
    {headers, options} = Keyword.pop!(options, :headers)
    {body, options} = Keyword.pop!(options, :body)
    Keyword.validate!(options, [:body_digest])

    datetime = DateTime.truncate(datetime, :second)
    datetime_string = DateTime.to_iso8601(datetime, :basic)
    date_string = Date.to_iso8601(datetime, :basic)
    url = URI.parse(url)
    body_digest = options[:body_digest] || hex(sha256(body))
    service = to_string(service)

    method = method |> Atom.to_string() |> String.upcase()

    aws_headers = [
      {"x-amz-content-sha256", body_digest},
      {"x-amz-date", datetime_string}
    ]

    aws_headers =
      if security_token do
        aws_headers ++ [{"x-amz-security-token", security_token}]
      else
        aws_headers
      end

    canonical_headers = headers ++ aws_headers

    ## canonical_headers needs to be sorted for canonical_request construction
    canonical_headers = Enum.sort(canonical_headers)

    signed_headers =
      Enum.map_intersperse(
        Enum.sort(canonical_headers),
        ";",
        &String.downcase(elem(&1, 0), :ascii)
      )

    canonical_headers =
      Enum.map_intersperse(canonical_headers, "\n", fn {name, value} -> [name, ":", value] end)

    canonical_query = canonical_query(url.query)

    canonical_request =
      [
        method,
        "\n",
        url.path,
        "\n",
        canonical_query,
        "\n",
        canonical_headers,
        "\n",
        "\n",
        signed_headers,
        "\n",
        body_digest
      ]
      |> IO.iodata_to_binary()

    string_to_sign =
      [
        "AWS4-HMAC-SHA256",
        "\n",
        datetime_string,
        "\n",
        "#{date_string}/#{region}/#{service}/aws4_request",
        "\n",
        hex(sha256(canonical_request))
      ]
      |> IO.iodata_to_binary()

    signature =
      aws_sigv4(
        string_to_sign,
        date_string,
        region,
        service,
        secret_access_key
      )

    credential = "#{access_key_id}/#{date_string}/#{region}/#{service}/aws4_request"

    authorization =
      "AWS4-HMAC-SHA256 Credential=#{credential},SignedHeaders=#{signed_headers},Signature=#{signature}"

    [{"authorization", authorization}] ++ aws_headers ++ headers
  end

  defp canonical_query(nil),
    do: ""

  defp canonical_query(query) do
    for item <- String.split(query, "&", trim: true) do
      case String.split(item, "=") do
        [name, value] -> [name, "=", value]
        [name] -> [name, "="]
      end
    end
    |> Enum.sort()
    |> Enum.intersperse("&")
  end

  @doc """
  Create AWS Signature v4 URL.

  https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-query-string-auth.html
  """
  def aws_sigv4_url(options) do
    {access_key_id, options} = Keyword.pop!(options, :access_key_id)
    {secret_access_key, options} = Keyword.pop!(options, :secret_access_key)
    {region, options} = Keyword.pop!(options, :region)
    {service, options} = Keyword.pop!(options, :service)
    {datetime, options} = Keyword.pop!(options, :datetime)
    {method, options} = Keyword.pop!(options, :method)
    {url, options} = Keyword.pop!(options, :url)
    {expires, options} = Keyword.pop(options, :expires, 86400)
    {headers, options} = Keyword.pop(options, :headers, [])
    {query_params, options} = Keyword.pop(options, :query_params, [])

    [] = options

    datetime = DateTime.truncate(datetime, :second)
    datetime_string = DateTime.to_iso8601(datetime, :basic)
    date_string = Date.to_iso8601(datetime, :basic)
    url = URI.parse(url)
    service = to_string(service)

    combined_headers =
      [host_header(url) | headers]
      |> Enum.sort(&compare_pair/2)

    signed_headers =
      Enum.map_intersperse(
        Enum.sort(combined_headers),
        ";",
        &String.downcase(elem(&1, 0), :ascii)
      )

    path = url.path
    method = method |> Atom.to_string() |> String.upcase()

    canonical_headers =
      Enum.map_intersperse(combined_headers, "\n", fn {name, value} -> [name, ":", value] end)

    aws_query_params =
      [
        {"X-Amz-Algorithm", "AWS4-HMAC-SHA256"},
        {"X-Amz-Credential", "#{access_key_id}/#{date_string}/#{region}/#{service}/aws4_request"},
        {"X-Amz-Date", datetime_string},
        {"X-Amz-Expires", to_string(expires)},
        {"X-Amz-SignedHeaders", signed_headers}
      ]
      |> Enum.sort(&compare_pair/2)

    # Extract existing query parameters
    custom_query_params =
      Enum.sort(
        Enum.to_list(URI.decode_query(url.query || "")) ++ query_params,
        &compare_pair/2
      )

    # Build canonical query string for signing
    query_to_sign =
      (custom_query_params ++ aws_query_params)
      |> Enum.sort(&compare_pair/2)
      |> encode_query_params()

    canonical_request =
      [
        method,
        "\n",
        path,
        "\n",
        query_to_sign,
        "\n",
        canonical_headers,
        "\n",
        "\n",
        signed_headers,
        "\n",
        "UNSIGNED-PAYLOAD"
      ]
      |> IO.iodata_to_binary()

    string_to_sign =
      [
        "AWS4-HMAC-SHA256",
        "\n",
        datetime_string,
        "\n",
        "#{date_string}/#{region}/#{service}/aws4_request",
        "\n",
        hex(sha256(canonical_request))
      ]
      |> IO.iodata_to_binary()

    signature =
      aws_sigv4(
        string_to_sign,
        date_string,
        region,
        service,
        secret_access_key
      )

    final_query_string =
      case custom_query_params do
        [] ->
          encode_query_params(aws_query_params)

        _ ->
          encode_query_params(custom_query_params) <> "&" <> encode_query_params(aws_query_params)
      end

    # Append the signature to the URL's query string
    put_in(url.query, final_query_string <> "&X-Amz-Signature=#{signature}")
  end

  # when it's http or https, the host implied when signing
  defp host_header(%URI{host: host, port: port}) when port in [80, 443] do
    {"host", host}
  end

  defp host_header(%URI{host: host, port: port}) do
    {"host", "#{host}:#{port}"}
  end

  def aws_sigv4(
        string_to_sign,
        date_string,
        region,
        service,
        secret_access_key
      ) do
    signature =
      ["AWS4", secret_access_key]
      |> hmac(date_string)
      |> hmac(region)
      |> hmac(service)
      |> hmac("aws4_request")
      |> hmac(string_to_sign)
      |> hex()

    signature
  end

  defp hex(data) do
    Base.encode16(data, case: :lower)
  end

  defp sha256(data) do
    :crypto.hash(:sha256, data)
  end

  defp hmac(key, data) do
    :crypto.mac(:hmac, :sha256, key, data)
  end

  defp encode_query_params(params) do
    Enum.map_join(params, "&", &pair/1)
  end

  defp compare_pair({key, value1}, {key, value2}), do: value1 < value2
  defp compare_pair({key_1, _}, {key_2, _}), do: key_1 < key_2

  defp pair({k, v}) do
    URI.encode_www_form(Kernel.to_string(k)) <> "=" <> aws_encode_www_form(Kernel.to_string(v))
  end

  # is basically the same as URI.encode_www_form
  # but doesn't use %20 instead of "+"
  defp aws_encode_www_form(str) when is_binary(str) do
    import Bitwise

    for <<c <- str>>, into: "" do
      case URI.char_unreserved?(c) do
        true -> <<c>>
        false -> "%" <> char_hex(bsr(c, 4)) <> char_hex(band(c, 15))
      end
    end
  end

  defp char_hex(n) when n <= 9, do: <<n + ?0>>
  defp char_hex(n), do: <<n + ?A - 10>>
end