class Aws::Resources::Collection

def [](index)
  if @size
    @batches[0][index]
  else
    raise "unable to index into a lazy loaded collection"
  end
end

def batch_enum
  case @limit
  when 0 then []
  when nil then non_empty_batches
  else limited_batches
  end
end

def batches
  ::Enumerator.new do |y|
    batch_enum.each do |batch|
      y << self.class.new([batch], size: batch.size)
    end
  end
end

def each(&block)
  enum = ::Enumerator.new do |y|
    batch_enum.each do |batch|
      batch.each do |band|
        y.yield(band)
      end
    end
  end
  enum.each(&block) if block
  enum
end

def first(count = nil)
  if count
    items = limit(count).to_a
    self.class.new([items], size: items.size)
  else
    begin
      each.next
    rescue StopIteration
      nil
    end
  end
end

def initialize(batches, options = {})
  @batches = batches
  @limit = options[:limit]
  @size = options[:size]
end

def limit(limit)
  Collection.new(@batches, limit: limit)
end

def limited_batches
  ::Enumerator.new do |y|
    yielded = 0
    @batches.each do |batch|
      batch = batch.take(@limit - yielded)
      if batch.size > 0
        y.yield(batch)
        yielded += batch.size
      end
      break if yielded == @limit
    end
  end
end

def non_empty_batches
  ::Enumerator.new do |y|
    @batches.each do |batch|
      y.yield(batch) if batch.size > 0
    end
  end
end

def size
  @size
end