#!/usr/bin/env ruby

require 'rubygems'
$:.unshift(File.join(File.dirname(__FILE__), "..", "lib"))
require 'chef_metal'
require 'chef/rest'
require 'chef/application'
require 'chef/knife'
require 'chef/run_context'
require 'chef/server_api'
require 'chef_metal/action_handler'
require 'chef_metal/version'
require 'chef_metal/chef_machine_spec'

class ChefMetal::Application < Chef::Application

  # Mimic self_pipe sleep from Unicorn to capture signals safely
  SELF_PIPE = []

  option :config_file,
    :short => "-c CONFIG",
    :long  => "--config CONFIG",
    :description => "The configuration file to use"

  option :log_level,
    :short        => "-l LEVEL",
    :long         => "--log_level LEVEL",
    :description  => "Set the log level (debug, info, warn, error, fatal)",
    :proc         => lambda { |l| l.to_sym }

  option :log_location,
    :short        => "-L LOGLOCATION",
    :long         => "--logfile LOGLOCATION",
    :description  => "Set the log file location, defaults to STDOUT - recommended for daemonizing",
    :proc         => nil

  option :node_name,
    :short => "-N NODE_NAME",
    :long => "--node-name NODE_NAME",
    :description => "The node name for this client",
    :proc => nil

  option :chef_server_url,
    :short => "-S CHEFSERVERURL",
    :long => "--server CHEFSERVERURL",
    :description => "The chef server URL",
    :proc => nil

  option :client_key,
    :short        => "-k KEY_FILE",
    :long         => "--client_key KEY_FILE",
    :description  => "Set the client key file location",
    :proc         => nil

  option :local_mode,
    :short        => "-z",
    :long         => "--local-mode",
    :description  => "Point chef-client at local repository",
    :boolean      => true

  option :chef_repo_path,
    :long         => '--chef-repo-path=PATH',
    :description  => "Path to Chef repository"

  option :chef_zero_port,
    :long         => "--chef-zero-port PORT",
    :description  => "Port to start chef-zero on"

  option :read_only,
    :long         => "--[no-]read-only",
    :description  => "Promise that execution will not modify the machine (helps with Docker in particular)"

  option :stream,
    :long         => "--[no-]stream",
    :default      => true,
    :boolean      => true,
    :description  => "Whether to stream output from the machine (default: true)"

  option :timeout,
    :long         => "--timeout=TIMEOUT",
    :default      => 15*60,
    :description  => "Time to wait for program to execute, or 0 for no timeout (default: 15 minutes)"

  def reconfigure
    super

    Chef::Config.chef_server_url = config[:chef_server_url] if config.has_key? :chef_server_url

    Chef::Config.chef_repo_path = config[:chef_repo_path] if config.has_key? :chef_repo_path

    Chef::Config.local_mode = config[:local_mode] if config.has_key?(:local_mode)
    if Chef::Config.local_mode && !Chef::Config.has_key?(:cookbook_path) && !Chef::Config.has_key?(:chef_repo_path)
      Chef::Config.chef_repo_path = Chef::Config.find_chef_repo_path(Dir.pwd)
    end
    Chef::Config.chef_zero.port = config[:chef_zero_port] if config[:chef_zero_port]
  end

  def setup_application
  end

  def load_config_file
    if !config.has_key?(:config_file)
      require 'chef/knife'
      config[:config_file] = Chef::Knife.locate_config_file
    end
    super
  end

  def run_application
    exit_code = 0

    Cheffish.honor_local_mode do
      command = cli_arguments.shift
      case command
      when 'execute'
        connect_to_machines(cli_arguments.shift) do |machine|
          machine.execute(action_handler, cli_arguments.join(' '), :read_only => config[:read_only], :stream => config[:stream], :timeout => config[:timeout].to_f)
          puts result.stdout if result.stdout != '' && !config[:stream] && Chef::Config.log_level != :debug
          STDERR.puts result.stderr if result.stderr != '' && !config[:stream] && Chef::Config.log_level != :debug
          exit_code = result.exitstatus if result.exitstatus != 0
        end
      when 'destroy'
        each_current_machine(cli_arguments.shift) do |driver, specs_and_options|
          driver.destroy_machines(action_handler, specs_and_options, parallelizer)
        end
      when 'allocate'
        each_new_machine(cli_arguments.shift) do |driver, specs_and_options|
          driver.allocate_machines(action_handler, specs_and_options, parallelizer)
        end
      when 'ready'
        each_new_machine(cli_arguments.shift) do |driver, specs_and_options|
          driver.allocate_machines(action_handler, specs_and_options, parallelizer)
          driver.ready_machines(action_handler, specs_and_options, parallelizer)
        end
      when 'setup'
        each_new_machine(cli_arguments.shift) do |driver, specs_and_options|
          driver.allocate_machines(action_handler, specs_and_options, parallelizer)
          driver.ready_machines(action_handler, specs_and_options, parallelizer) do |machine|
            machine.setup_convergence(action_handler)
          end
        end
      when 'converge'
        each_new_machine(cli_arguments.shift) do |driver, specs_and_options|
          driver.allocate_machines(action_handler, specs_and_options, parallelizer)
          driver.ready_machines(action_handler, specs_and_options, parallelizer) do |machine|
            # TODO upload files?  Maybe they should be in machine_options?
            machine.setup_convergence(action_handler)
            machine.converge(action_handler)
          end
        end
      when 'reconverge'
        connect_to_machines(cli_arguments.shift) do |machine|
          machine.converge(action_handler)
        end
      when 'stop'
        each_current_machine(cli_arguments.shift) do |driver, specs_and_options|
          driver.stop_machines(action_handler, specs_and_options, parallelizer)
        end
      when 'cat'
        connect_to_machines(cli_arguments.shift) do |machine|
          cli_arguments.each do |remote_path|
            puts machine.read_file(remote_path)
          end
        end
      when 'cp'
        machines = {}
        to = cli_arguments.pop
        if to =~ /^([^\/:]+):(.+)$/
          to_server = $1
          machines[to_server] ||= ChefMetal.connect_to_machine(to_server)
          to_path = $2
          to_is_directory = machines[to_server].is_directory?(to_path)
        else
          to_server = nil
          to_path = File.absolute_path(to)
        end

        cli_arguments.each do |from|
          if from =~ /^([^\/:]+):(.+)$/
            from_server = $1
            from_path = $2
            if to_server
              raise "Cannot copy from one server to another, or intraserver (from=#{from}, to=#{to})"
            end

            machines[from_server] ||= connect_to_machine(from_server)
            if File.directory?(to_path)
              machines[from_server].download_file(action_handler, from_path, "#{to_path}/#{File.basename(from_path)}")
            else
              machines[from_server].download_file(action_handler, from_path, to_path)
            end
          else
            from_server = nil
            from_path = File.absolute_path(from)
            if !to_server
              raise "Cannot copy two local files.  One of the arguments must be MACHINE:PATH.  (from=#{from}, to=#{to})"
            end

            if to_is_directory
              machines[to_server].upload_file(action_handler, from_path, "#{to_path}/#{File.basename(from_path)}")
            else
              machines[to_server].upload_file(action_handler, from_path, to_path)
            end
          end
        end
      else
        Chef::Log.error("Command '#{command}' unrecognized")
      end
    end

    exit(exit_code) if exit_code != 0
  end

  private

  def rest
    @rest ||= Chef::ServerAPI.new()
  end

  def machine_specs(*specs)
    names = specs.collect_concat { |spec| spec.split(',') }.uniq
    parallelizer.parallelize(names) { |name| ChefMetal::ChefMachineSpec.get(name) }.to_a
  end

  def each_new_machine(spec)
    driver = ChefMetal.default_driver
    specs_and_options = {}
    machine_specs(spec).each do |machine_spec|
      specs_and_options[machine_spec] = driver.config[:machine_options]
    end
    [ [ driver,  specs_and_options ] ]
  end

  def each_current_machine(spec)
    grouped = machine_specs(spec).group_by { |machine_spec| machine_spec.driver_url }
    parallelizer.parallelize(grouped) do |driver_url, machine_specs|
      if driver_url
        driver = ChefMetal.driver_for_url(driver_url)
        specs_and_options = {}
        machine_specs(spec).each do |machine_spec|
          specs_and_options[machine_spec] = driver.config[:machine_options]
        end
        yield driver, specs_and_options
      end
    end.to_a
  end

  def connect_to_machines(spec)
    grouped = machine_specs(spec).group_by { |machine_spec| machine_spec.driver_url }
    grouped.collect_concat do |driver_url, machine_specs|
      machine_specs.map { |machine_spec| ChefMetal.connect_to_machine(machine_spec) }
    end
  end

  def parallelizer
    Chef::ChefFS::Parallelizer
  end

  def action_handler
    @action_handler ||= ActionHandler.new
  end

  class ActionHandler < ChefMetal::ActionHandler
    def recipe_context
      # TODO: somehow remove this code; should context really always be needed?
      node = Chef::Node.new
      node.name 'nothing'
      node.automatic[:platform] = 'metal'
      node.automatic[:platform_version] = ChefMetal::VERSION
      Chef::RunContext.new(node, {},
        Chef::EventDispatch::Dispatcher.new(Chef::Formatters::Doc.new(STDOUT,STDERR)))
    end
  end
end

ChefMetal::Application.new.run