class JSON::JWK::Set

def [](kid)
  detect do |jwk|
    jwk[:kid] && jwk[:kid] == kid
  end
end

def as_json(options = {})
  # NOTE: Array.new wrapper is requied to avoid CircularReferenceError
  {keys: Array.new(self)}
end

def content_type
  'application/jwk-set+json'
end

def initialize(*jwks)
  jwks = if jwks.first.is_a?(Hash) && (keys = jwks.first.with_indifferent_access[:keys])
    keys
  else
    jwks
  end
  jwks = Array(jwks).flatten.collect do |jwk|
    JWK.new jwk
  end
  replace jwks
end