gnat

defmodule Gnat.Jetstream.API.Object do
  @moduledoc """
  API for interacting with the JetStream Object Store

  Learn more about Object Store: https://docs.nats.io/nats-concepts/jetstream/obj_store
  """
  alias Gnat.Jetstream.API.{Consumer, Stream, Util}
  alias Gnat.Jetstream.API.Object.Meta

  @stream_prefix "OBJ_"
  @subject_prefix "$O."

  @type bucket_opt ::
          {:description, String.t()}
          | {:max_bucket_size, integer()}
          | {:max_chunk_size, integer()}
          | {:placement, Stream.placement()}
          | {:replicas, non_neg_integer()}
          | {:storage, :file | :memory}
          | {:ttl, non_neg_integer()}
  @spec create_bucket(Gnat.t(), String.t(), list(bucket_opt)) ::
          {:ok, Stream.info()} | {:error, any()}
  def create_bucket(conn, bucket_name, params \\ []) do
    with :ok <- validate_bucket_name(bucket_name) do
      stream = %Stream{
        name: stream_name(bucket_name),
        subjects: stream_subjects(bucket_name),
        description: Keyword.get(params, :description),
        discard: :new,
        allow_rollup_hdrs: true,
        max_age: Keyword.get(params, :ttl, 0),
        max_bytes: Keyword.get(params, :max_bucket_size, -1),
        max_msg_size: Keyword.get(params, :max_chunk_size, -1),
        num_replicas: Keyword.get(params, :replicas, 1),
        storage: Keyword.get(params, :storage, :file),
        placement: Keyword.get(params, :placement),
        duplicate_window: adjust_duplicate_window(Keyword.get(params, :ttl, 0))
      }

      Stream.create(conn, stream)
    end
  end

  @spec delete_bucket(Gnat.t(), String.t()) :: :ok | {:error, any}
  def delete_bucket(conn, bucket_name) do
    Stream.delete(conn, stream_name(bucket_name))
  end

  @spec delete(Gnat.t(), String.t(), String.t()) :: :ok | {:error, any}
  def delete(conn, bucket_name, object_name) do
    with {:ok, meta} <- info(conn, bucket_name, object_name),
         meta <- %Meta{meta | deleted: true},
         topic <- meta_stream_topic(bucket_name, object_name),
         {:ok, body} <- Jason.encode(meta),
         {:ok, _msg} <- Gnat.request(conn, topic, body, headers: [{"Nats-Rollup", "sub"}]) do
      filter = chunk_stream_topic(meta)
      Stream.purge(conn, stream_name(bucket_name), nil, %{filter: filter})
    end
  end

  @spec get(Gnat.t(), String.t(), String.t(), (binary -> any())) :: :ok | {:error, any}
  def get(conn, bucket_name, object_name, chunk_fun) do
    with {:ok, %{config: _stream}} <- Stream.info(conn, stream_name(bucket_name)),
         {:ok, meta} <- info(conn, bucket_name, object_name) do
      receive_chunks(conn, meta, chunk_fun)
    end
  end

  @spec info(Gnat.t(), String.t(), String.t()) :: {:ok, Meta.t()} | {:error, any}
  def info(conn, bucket_name, object_name) do
    with {:ok, _stream_info} <- Stream.info(conn, stream_name(bucket_name)) do
      Stream.get_message(conn, stream_name(bucket_name), %{
        last_by_subj: meta_stream_topic(bucket_name, object_name)
      })
      |> case do
        {:ok, message} ->
          meta = json_to_meta(message.data)
          {:ok, meta}

        error ->
          error
      end
    end
  end

  @type list_option :: {:show_deleted, boolean()}
  @spec list(Gnat.t(), String.t(), list(list_option())) :: {:error, any} | {:ok, list(Meta.t())}
  def list(conn, bucket_name, options \\ []) do
    with {:ok, %{config: stream}} <- Stream.info(conn, stream_name(bucket_name)),
         topic <- Util.reply_inbox(),
         {:ok, sub} <- Gnat.sub(conn, self(), topic),
         {:ok, consumer} <-
           Consumer.create(conn, %Consumer{
             stream_name: stream.name,
             deliver_subject: topic,
             deliver_policy: :last_per_subject,
             filter_subject: meta_stream_subject(bucket_name),
             ack_policy: :none,
             max_ack_pending: nil,
             replay_policy: :instant,
             max_deliver: 1
           }),
         {:ok, messages} <- receive_all_metas(sub, consumer.num_pending) do
      :ok = Gnat.unsub(conn, sub)
      :ok = Consumer.delete(conn, stream.name, consumer.name)

      show_deleted = Keyword.get(options, :show_deleted, false)

      if show_deleted do
        {:ok, messages}
      else
        {:ok, Enum.reject(messages, &(&1.deleted == true))}
      end
    end
  end

  @spec put(Gnat.t(), String.t(), String.t(), File.io_device()) ::
          {:ok, Meta.t()} | {:error, any()}
  def put(conn, bucket_name, object_name, io) do
    nuid = Util.nuid()
    chunk_topic = chunk_stream_topic(bucket_name, nuid)

    with {:ok, %{config: _}} <- Stream.info(conn, stream_name(bucket_name)),
         :ok <- purge_prior_chunks(conn, bucket_name, object_name),
         {:ok, chunks, size, digest} <- send_chunks(conn, io, chunk_topic) do
      object_meta = %Meta{
        name: object_name,
        bucket: bucket_name,
        nuid: nuid,
        size: size,
        chunks: chunks,
        digest: "SHA-256=#{Base.url_encode64(digest)}"
      }

      topic = meta_stream_topic(bucket_name, object_name)
      body = Jason.encode!(object_meta)

      case Gnat.request(conn, topic, body, headers: [{"Nats-Rollup", "sub"}]) do
        {:ok, _} ->
          {:ok, object_meta}

        error ->
          error
      end
    end
  end

  defp stream_name(bucket_name) do
    "#{@stream_prefix}#{bucket_name}"
  end

  defp stream_subjects(bucket_name) do
    [
      chunk_stream_subject(bucket_name),
      meta_stream_subject(bucket_name)
    ]
  end

  defp chunk_stream_subject(bucket_name) do
    "#{@subject_prefix}#{bucket_name}.C.>"
  end

  defp chunk_stream_topic(bucket_name, nuid) do
    "#{@subject_prefix}#{bucket_name}.C.#{nuid}"
  end

  defp chunk_stream_topic(%Meta{bucket: bucket, nuid: nuid}) do
    "#{@subject_prefix}#{bucket}.C.#{nuid}"
  end

  defp meta_stream_subject(bucket_name) do
    "#{@subject_prefix}#{bucket_name}.M.>"
  end

  defp meta_stream_topic(bucket_name, object_name) do
    key = Base.url_encode64(object_name)
    "#{@subject_prefix}#{bucket_name}.M.#{key}"
  end

  @two_minutes_in_nanoseconds 1_200_000_000
  # The `duplicate_window` can't be greater than the `max_age`. The default `duplicate_window`
  # is 2 minutes. We'll keep the 2 minute window UNLESS the ttl is less than 2 minutes
  defp adjust_duplicate_window(ttl) when ttl > 0 and ttl < @two_minutes_in_nanoseconds, do: ttl
  defp adjust_duplicate_window(_ttl), do: @two_minutes_in_nanoseconds

  defp json_to_meta(json) do
    raw = Jason.decode!(json)

    %{
      "bucket" => bucket,
      "chunks" => chunks,
      "digest" => digest,
      "name" => name,
      "nuid" => nuid,
      "size" => size
    } = raw

    %Meta{
      bucket: bucket,
      chunks: chunks,
      digest: digest,
      deleted: Map.get(raw, "deleted", false),
      name: name,
      nuid: nuid,
      size: size
    }
  end

  defp purge_prior_chunks(conn, bucket, name) do
    case info(conn, bucket, name) do
      {:ok, meta} ->
        Stream.purge(conn, stream_name(bucket), nil, %{filter: chunk_stream_topic(meta)})

      {:error, %{"code" => 404}} ->
        :ok

      {:error, other} ->
        {:error, other}
    end
  end

  defp receive_all_metas(sid, num_pending, messages \\ [])

  defp receive_all_metas(_sid, 0, messages) do
    {:ok, messages}
  end

  defp receive_all_metas(sid, remaining, messages) do
    receive do
      {:msg, %{sid: ^sid, body: body}} ->
        meta = json_to_meta(body)
        receive_all_metas(sid, remaining - 1, [meta | messages])
    after
      10_000 ->
        {:error, :timeout_waiting_for_messages}
    end
  end

  defp receive_chunks(conn, %Meta{} = meta, chunk_fun) do
    topic = chunk_stream_topic(meta)
    stream = stream_name(meta.bucket)
    inbox = Util.reply_inbox()
    {:ok, sub} = Gnat.sub(conn, self(), inbox)

    {:ok, consumer} =
      Consumer.create(conn, %Consumer{
        stream_name: stream,
        deliver_subject: inbox,
        deliver_policy: :all,
        filter_subject: topic,
        ack_policy: :none,
        max_ack_pending: nil,
        replay_policy: :instant,
        max_deliver: 1
      })

    :ok = receive_chunks(sub, meta.chunks, chunk_fun)

    :ok = Gnat.unsub(conn, sub)
    :ok = Consumer.delete(conn, stream, consumer.name)
  end

  defp receive_chunks(_sub, 0, _chunk_fun) do
    :ok
  end

  defp receive_chunks(sub, remaining, chunk_fun) do
    receive do
      {:msg, %{sid: ^sub, body: body}} ->
        chunk_fun.(body)
        receive_chunks(sub, remaining - 1, chunk_fun)
    after
      10_000 ->
        {:error, :timeout_waiting_for_messages}
    end
  end

  @chunk_size 128 * 1024
  defp send_chunks(conn, io, topic) do
    sha = :crypto.hash_init(:sha256)
    size = 0
    chunks = 0
    send_chunks(conn, io, topic, sha, size, chunks)
  end

  defp send_chunks(conn, io, topic, sha, size, chunks) do
    case IO.binread(io, @chunk_size) do
      :eof ->
        sha = :crypto.hash_final(sha)
        {:ok, chunks, size, sha}

      {:error, err} ->
        {:error, err}

      bytes ->
        sha = :crypto.hash_update(sha, bytes)
        size = size + byte_size(bytes)
        chunks = chunks + 1

        case Gnat.request(conn, topic, bytes) do
          {:ok, _} ->
            send_chunks(conn, io, topic, sha, size, chunks)

          error ->
            error
        end
    end
  end

  defp validate_bucket_name(name) do
    case Regex.match?(~r/^[a-zA-Z0-9_-]+$/, name) do
      true -> :ok
      false -> {:error, "invalid bucket name"}
    end
  end
end