class Nokogiri::XML::NodeSet

def << node
  push(node)
end

def add_class name
  each do |el|
    next unless el.respond_to? :get_attribute
    classes = el.get_attribute('class').to_s.split(" ")
    el.set_attribute('class', classes.push(name).uniq.join(" "))
  end
  self
end

def after datum
  last.after datum
end

def at path, ns = {}
  return self[path] if path.is_a?(Numeric)
  search(path, ns).first
end

def attr key, value = nil, &blk
  if value or blk
    each do |el|
      el.set_attribute(key, value || blk[el])
    end
    return self      
  end    
  if key.is_a? Hash
    key.each { |k,v| self.attr(k,v) }
    return self
  else
    return self[0].get_attribute(key)
  end
end

def before datum
  first.before datum
end

def each(&block)
  x = 0
  while x < length
    yield self[x]
    x += 1
  end
end

def empty?
  length == 0
end

def first
  self[0]
end

def initialize document, list = []
  @document = document
  list.each { |x| self << x }
  yield self if block_given?
end

def inner_html
  collect{|j| j.inner_html}.join('')
end

def inner_text
  collect{|j| j.inner_text}.join('')
end

def last
  self[length - 1]
end

def remove_attr name
  each do |el|
    next unless el.respond_to? :remove_attribute
    el.remove_attribute(name)
  end
  self      
end

def remove_class name = nil
  each do |el|
    next unless el.respond_to? :get_attribute
    if name
      classes = el.get_attribute('class').to_s.split(" ")
      el.set_attribute('class', (classes - [name]).uniq.join(" "))
    else
      el.remove_attribute("class")
    end
  end
  self
end

def search *paths
  sub_set = NodeSet.new(document)
  document.decorate(sub_set)
  each do |node|
    node.search(*paths).each do |sub_node|
      sub_set << sub_node
    end
  end
  sub_set
end

def size
  length
end

def to_ary
  to_a
end

def to_html
  map { |x| x.to_html }.join('')
end

def to_s
  map { |x| x.to_s }.join
end

def to_xml *args
  map { |x| x.to_xml(*args) }.join('')
end

def unlink
  each { |node| node.unlink }
  self
end

def wrap(html, &blk)
  each do |j|
    new_parent = Nokogiri.make(html, &blk)
    j.replace(new_parent)
    nest = new_parent
    if nest.child
      nest = nest.child until nest.child.nil?
    end
    j.parent = nest
  end
  self
end