require 'nokogiri'

# Usage:
#
#   use Rack::Injector do |inject|
#     inject.before 'h1', '<div class="injected">This was injected before</div>'
#     inject.after 'h1', '<div class="injected">This was injected after</div>'
#   end
class Rack::Injector
  def initialize(app)
    @app = app
    @afters = {}
    @befores = {}
    yield self if block_given?
  end

  def call(env)
    status, headers, body = @app.call(env)
    response = Rack::Response.new(inject!(body.join), status, headers)
    response.finish
  end

  def before(selector, content)
    @befores[selector] ||= []
    @befores[selector] << content
  end

  def after(selector, content)
    @afters[selector] ||= []
    @afters[selector] << content
  end

  def inject!(body)
    doc = Nokogiri::HTML(body)

    @afters.each do |selector, contents|
      doc.search(selector).each do |node|
        node.after(contents.join("\n")) if node
      end
    end

    @befores.each do |selector, contents|
      doc.search(selector).each do |node|
        node.before(contents.join("\n")) if node
      end
    end

    [doc.to_html]
  end
end