#!/usr/bin/env ruby

require 'rubygems'
require 'chef'
require 'chef/node'
require 'chef/rest'
require 'optparse'
require 'ostruct'

Chef::Config.from_file("/etc/chef/server.rb")
API_USERNAME=ENV['CHEF_USERNAME']
API_PASSWORD=ENV['CHEF_PASSWORD']
@rest = Chef::REST.new(Chef::Config[:registration_url])
@rest.authenticate(API_USERNAME, API_PASSWORD)

def get_nodes_with_webapp(webapp)
  distinct = {}
  results = @rest.get_rest("search/node?q=#{webapp}&a=recipes")
  results.each do |result|
    distinct[result['id']] = 1
  end
  rv = []
  distinct.each do |node|
    node_name = node[0].sub('node_','')
    node_name.gsub!(".", "_")
    begin
      fullnode = @rest.get_rest("nodes/#{node_name}")
      if fullnode.fqdn
        rv << node_name
      end
    rescue
    end
  end
  rv
end

def update_webapps(webapp, update_type)
  node_names = get_nodes_with_webapp(webapp)
  node_names.each do |node_name|
    node = @rest.get_rest("nodes/#{node_name}")
    # for each webapp under management on each node
    # there's an attribute set "#{webapp}_deploy_action"
    # that's normally set to "nothing" -- here we're
    # resetting it "deploy" or "rollback"
    attr = webapp + "_deploy_action"
    node[attr] = update_type
    puts "Storing node data for #{node_name}..."
    begin
      retries = 5
      @rest.put_rest("nodes/#{node_name}", node)
    rescue Net::HTTPFatalError
      retry if (retries -= 1) > 0
    end
    puts "Done."
  end
end


if __FILE__ == $0
  options = OpenStruct.new
  opts = OptionParser.new do |opts|
    opts.on("-w", "--webapp [ARG]", "Web app to update") do |w|
      if w.nil?
        puts "--webapp cannot be null"
        exit
      else
        options.webapp = w
      end
    end
    opts.on("-u", "--update [ARG]", "Type of update (deploy|rollback)") do |u|
      if u.nil?
        puts "--update must be 'deploy' or 'rollback'"
        exit
      else
        options.update = u
      end
    end
  end
  begin
    opts.parse!(ARGV)
  rescue OptionParser::ParseError => e
    puts e
    # put a usage message here
  end
  if options.webapp.nil? || options.update.nil?
    puts "Usage: #{$0} --webapp <webapp> --update <update>"
    exit
  end
  update_webapps(options.webapp, options.update)
end