require 'rubygems'
require 'digest/md5'
require 'json'

class BitField

	attr_reader :buckets

	def initialize (arg)
		@bits = arg['bits'] ? arg['bits'] : 307
		@bucketsize = arg['bucketsize'] ? arg['bucketsize'] : 31

		if arg['buckets']
			@buckets = arg['buckets']
		else
			@buckets = Array.new((@bits / @bucketsize), 0)
		end
	end

	def setbit (bit)
		b = bit % @bucketsize
		@buckets[(bit / @bucketsize)] |= (1 << b)
	end

	def bitset? (bit)
		b = bit % @bucketsize
		return (@buckets[(bit / @bucketsize)] & (1 << b) == 0) ? false : true
	end
end

class BloomFilter

	def initialize (arg)
		throw 'Missing salts key' unless arg['salts']
		@salts = arg['salts']
		@bits = arg['bits'] || 200
		@bucketsize = arg['bucketsize'] || 31

		bfarg = { 'bits' => @bits, 'bucketsize' => @bucketsize }
		bfarg['buckets'] = arg['buckets'] if arg['buckets']
		@bitfield = BitField.new(bfarg)
	end

	def hashes (v)
		r = []
		@salts.each { |salt|
			r.push(
				Digest::MD5.hexdigest("#{salt}#{v}")[0..7].to_i(16) % @bits
			)
		}
		return r
	end

	def add (v)
		self.hashes(v).each { |bit| @bitfield.setbit(bit) }
	end

	def test (v)
		self.hashes(v).each { |bit|
			next if @bitfield.bitset?(bit)
			return false
		}
		return true
	end

	def to_json
		return {
			'buckets' => @bitfield.buckets,
			'bucketsize' => @bucketsize,
			'bits' => @bits,
			'salts' => @salts
		}.to_json
	end

end