class Nokogiri::XML::NodeSet

def == other
  return false unless other.is_a?(Nokogiri::XML::NodeSet)
  return false unless length == other.length
  each_with_index do |node, i|
    return false unless node == other[i]
  end
  true
end

def > selector
  ns = document.root.namespaces
  xpath CSS.xpath_for(selector, :prefix => "./", :ns => ns).first
end

def add_class name
  each do |el|
    classes = el['class'].to_s.split(/\s+/)
    el['class'] = classes.push(name).uniq.join " "
  end
  self
end

def after datum
  last.after datum
end

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

def at_css *rules
  css(*rules).first
end

def at_xpath *paths
  xpath(*paths).first
end

def attr key, value = nil, &blk
  unless Hash === key || key && (value || blk)
    return first.attribute(key)
  end
  hash = key.is_a?(Hash) ? key : { key => value }
  hash.each { |k,v| each { |el| el[k] = v || blk[el] } }
  self
end

def before datum
  first.before datum
end

def children
  inject(NodeSet.new(document)) { |set, node| set += node.children }
end

def css *paths
  handler = ![
    Hash, String, Symbol
  ].include?(paths.last.class) ? paths.pop : nil
  ns = paths.last.is_a?(Hash) ? paths.pop : nil
  sub_set = NodeSet.new(document)
  each do |node|
    doc = node.document
    search_ns = ns || (doc.root ? doc.root.namespaces : {})
    xpaths = paths.map { |rule|
      [
        CSS.xpath_for(rule.to_s, :prefix => ".//", :ns => search_ns),
        CSS.xpath_for(rule.to_s, :prefix => "self::", :ns => search_ns)
      ].join(' | ')
    }
    sub_set += node.xpath(*(xpaths + [search_ns, handler].compact))
  end
  document.decorate(sub_set)
  sub_set
end

def each(&block)
  0.upto(length - 1) do |x|
    yield self[x]
  end
end

def empty?
  length == 0
end

def filter expr
  find_all { |node| node.matches?(expr) }
end

def first n = nil
  return self[0] unless n
  list = []
  0.upto(n - 1) do |i|
    list << self[i]
  end
  list
end

def index(node)
  each_with_index { |member, j| return j if member == node }
  nil
end

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

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

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

def inspect
  "[#{map { |c| c.inspect }.join ', '}]"
end

def last
  self[length - 1]
end

def pop
  return nil if length == 0
  delete last
end

def remove_attr name
  each { |el| el.delete name }
  self
end

def remove_class name = nil
  each do |el|
    if name
      classes = el['class'].to_s.split(/\s+/)
      if classes.empty?
        el.delete 'class'
      else
        el['class'] = (classes - [name]).uniq.join " "
      end
    else
      el.delete "class"
    end
  end
  self
end

def reverse
  node_set = NodeSet.new(document)
  (length - 1).downto(0) do |x|
    node_set.push self[x]
  end
  node_set
end

def search *paths
  handler = ![
    Hash, String, Symbol
  ].include?(paths.last.class) ? paths.pop : nil
  ns = paths.last.is_a?(Hash) ? paths.pop : nil
  sub_set = NodeSet.new(document)
  paths.each do |path|
    sub_set += send(
      path =~ /^(\.\/|\/)/ ? :xpath : :css,
      *(paths + [ns, handler]).compact
    )
  end
  document.decorate(sub_set)
  sub_set
end

def shift
  return nil if length == 0
  delete first
end

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

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

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

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

def wrap(html, &blk)
  each do |j|
    new_parent = document.parse(html).first
    j.add_next_sibling(new_parent)
    new_parent.add_child(j)
  end
  self
end

def xpath *paths
  handler = ![
    Hash, String, Symbol
  ].include?(paths.last.class) ? paths.pop : nil
  ns = paths.last.is_a?(Hash) ? paths.pop : nil
  sub_set = NodeSet.new(document)
  each do |node|
    sub_set += node.xpath(*(paths + [ns, handler].compact))
  end
  document.decorate(sub_set)
  sub_set
end