class Nokogiri::XML::NodeSet

def self.new document, list = [] # :nodoc:
  set = NodeSet.wrap(LibXML.xmlXPathNodeSetCreate(nil))
  set.document = document
  list.each { |x| set << x }
  yield set if block_given?
  set
end

def self.wrap(ptr) # :nodoc:
  set = allocate
  set.cstruct = LibXML::XmlNodeSet.new(ptr)
  set
end

def &(node_set) # :nodoc:
  raise(ArgumentError, "node_set must be a Nokogiri::XML::NodeSet") unless node_set.is_a?(XML::NodeSet)
  new_set_ptr = LibXML.xmlXPathIntersection(cstruct, node_set.cstruct)
  NodeSet.wrap(new_set_ptr)
end

def +(node_set) # :nodoc:
  raise(ArgumentError, "node_set must be a Nokogiri::XML::NodeSet") unless node_set.is_a?(XML::NodeSet)
  new_set_ptr = LibXML::xmlXPathNodeSetMerge(nil, self.cstruct)
  new_set_ptr = LibXML::xmlXPathNodeSetMerge(new_set_ptr, node_set.cstruct)
  NodeSet.wrap(new_set_ptr)
end

def -(node_set) # :nodoc:
  raise(ArgumentError, "node_set must be a Nokogiri::XML::NodeSet") unless node_set.is_a?(XML::NodeSet)
  new_set_ptr = LibXML.xmlXPathNodeSetMerge(nil, self.cstruct)
  other_nodetab = node_set.cstruct.nodeTab
  node_set.cstruct[:nodeNr].times do |j|
    LibXML.xmlXPathNodeSetDel(new_set_ptr, other_nodetab[j])
  end        
  NodeSet.wrap(new_set_ptr)
end

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 [](*args) # :nodoc:
  raise(ArgumentError, "got #{args.length} arguments, expected 1 (or 2)") if args.length > 2
  if args.length == 2
    beg = args[0]
    len = args[1]
    beg += cstruct[:nodeNr] if beg < 0
    return subseq(beg, len)
  end
  arg = args[0]
  return subseq(arg.first, arg.last-arg.first+1) if arg.is_a?(Range)
  index_at(arg)
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 = document.root ? document.root.namespaces : {}
  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 delete(node) # :nodoc:
  raise(ArgumentError, "node must be a Nokogiri::XML::Node") unless node.is_a?(XML::Node)
  if LibXML.xmlXPathNodeSetContains(cstruct, node.cstruct) != 0
    LibXML.xmlXPathNodeSetDel(cstruct, node.cstruct)
    return node
  end
  return nil
end

def dup # :nodoc:
  dup = LibXML.xmlXPathNodeSetMerge(nil, self.cstruct)
  NodeSet.wrap(dup)
end

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

def empty?
  length == 0
end

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

def include?(node) # :nodoc:
  raise(ArgumentError, "node must be a Nokogiri::XML::Node") unless node.is_a?(XML::Node)
  (LibXML.xmlXPathNodeSetContains(cstruct, node.cstruct) != 0) ? true : false
end

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

def index_at(number) # :nodoc:
  return nil if (number >= cstruct[:nodeNr] || number.abs > cstruct[:nodeNr])
  number = number + cstruct[:nodeNr] if number < 0
  Node.wrap(cstruct.nodeAt(number))
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 length # :nodoc:
  cstruct.pointer.null? ? 0 : cstruct[:nodeNr]
end

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

def push(node) # :nodoc:
  raise(ArgumentError, "node must be a Nokogiri::XML::Node") unless node.is_a?(XML::Node)
  LibXML.xmlXPathNodeSetAdd(cstruct, node.cstruct)
  self
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
  ns = paths.last.is_a?(Hash) ? paths.pop : document.root.namespaces
  sub_set = NodeSet.new(document)
  document.decorate(sub_set)
  each do |node|
    node.search(*(paths + [ns])).each do |sub_node|
      sub_set << sub_node
    end
  end
  sub_set
end

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

def subseq(beg, len) # :nodoc:
  return nil if beg > cstruct[:nodeNr]
  return nil if beg < 0 || len < 0
  set = NodeSet.wrap(LibXML.xmlXPathNodeSetCreate(nil))
  beg.upto(beg+len-1) do |j|
    LibXML.xmlXPathNodeSetAdd(set.cstruct, cstruct.nodeAt(j));
  end
  set
end

def to_a # :nodoc:
  cstruct.nodeTab.collect { |node| Node.wrap(node) }
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 unlink # :nodoc:
  # TODO: is this simpler implementation viable:
  #  cstruct.nodeTab.collect {|node| Node.wrap(node)}.each(&:unlink)
  # ?
  nodetab = cstruct.nodeTab
  cstruct[:nodeNr].times do |j|
    node = Node.wrap(nodetab[j])
    node.unlink
    nodetab[j] = node.cstruct.pointer
  end
  cstruct.nodeTab = nodetab
  self
end

def wrap(html, &blk)
  each do |j|
    new_parent = Nokogiri.make(html, &blk)
    j.parent.add_child(new_parent)
    new_parent.add_child(j)
  end
  self
end