defmodule AgentStorage do
  use Agent

  def start_link(initial_value \\ %{}) do
    Agent.start_link(fn -> %{} end)
  end

  def open do
    {:ok, pid} = start_link()
    pid
  end

  def put(pid, key, value) do
    Agent.update(pid, fn state -> Map.put(state, key, value) end)
    pid
  end

  def get(pid, key) do
    Agent.get(pid, fn state -> Map.get(state, key) end)
  end

  def close(pid) do
    Agent.stop(pid)
  end
end

defmodule EtsStorage do
  def open do
    :ets.new(:lookup, [:set, :protected])
  end

  def put(id, key, value) do
    :ets.insert(id, {key, value})
    id
  end

  def get(id, key) do
    :ets.lookup_element(id, key, 2)
  end

  def close(id) do
    :ets.delete(id)
  end
end

defmodule Suite do
  def suite(impl, new, close \\ fn _ -> :ok end) do
    db = new.()

    db = Enum.reduce(0..10000, db, fn i, acc ->
      impl.put(acc, "key#{i}", "value#{i}")
    end)

    "value5050" =
      db
      |> impl.get("key5050")

    close.(db)
  end
end

map_suite = fn ->
  Suite.suite(Map, &Map.new/0)
end

memory_storage_suite = fn ->
  Suite.suite(MemoryStorage.Database, &MemoryStorage.Database.open/0)
end

agent_suite = fn ->
  Suite.suite(AgentStorage, &AgentStorage.open/0, &AgentStorage.close/1)
end

ets_suite = fn ->
  Suite.suite(EtsStorage, &EtsStorage.open/0, &EtsStorage.close/1)
end

Benchee.run(%{
  # "map" => map_suite,
  "memory_storage" => memory_storage_suite,
  "agent" => agent_suite,
  "ets" => ets_suite,
})
