#!/usr/bin/env ruby

require 'thor'
require 'cucumber-chef'


# $logger = Cucumber::Chef.logger

class CucumberChef < Thor
  include Thor::Actions

  no_tasks do

    def initalize_config
      source_dir = File.expand_path(File.join(File.dirname(__FILE__), "..", "lib", "cucumber", "chef", "templates", "cucumber-chef"))
      destination_dir = File.expand_path(File.join(Cucumber::Chef.locate_parent(".chef"), ".cucumber-chef"))
      FileUtils.mkdir_p(destination_dir)

      CucumberChef.source_root(source_dir)

      templates = {
        "config-rb.erb" => "config.rb"
      }

      templates.each do |source, destination|
        template(source, File.join(destination_dir, destination))
      end
      puts
      say "Ucanhaz Cucumber-Chef now! Rock on.", :green
    end

    def create_project(project)
      @project = project
      source_dir = File.expand_path(File.join(File.dirname(__FILE__), "..", "lib", "cucumber", "chef", "templates", "cucumber"))
      destination_dir = Cucumber::Chef.locate_parent(".chef")

      CucumberChef.source_root source_dir
      templates = {
        "readme.erb" => "features/#{project}/README",
        "example_feature.erb" => "features/#{project}/#{project}.feature",
        "example_steps.erb" => "features/#{project}/step_definitions/#{project}_steps.rb",
        "env.rb" => "features/support/env.rb",
        "cc-hooks.rb" => "features/support/cc-hooks.rb",
        "readme-data_bags.erb" => "features/support/data_bags/README",
        "readme-roles.erb" => "features/support/roles/README",
        "readme-keys.erb" => "features/support/keys/README",
        "readme-environments.erb" => "features/support/environments/README"
      }

      templates.each do |source, destination|
        template(source, File.join(destination_dir, destination))
      end
    end

    def boot
      tag = Cucumber::Chef.tag("cucumber-chef")
      puts(tag)
      Cucumber::Chef.boot(tag)
      $logger = Cucumber::Chef.logger

      @is_rc = Cucumber::Chef.is_rc?

      @options.test? and Cucumber::Chef::Config.test
    end

    def fatal(message)
      puts(set_color(message, :red, :bold))
      exit(255)
    end

  end

################################################################################

  desc "init", "Initalize cucumber-chef configuration"
  def init
    initalize_config
  end

################################################################################
# SETUP
################################################################################

  desc "setup", "Setup the cucumber-chef test lab"
  method_option :test, :type => :boolean, :desc => "INTERNAL USE ONLY"
  def setup
    boot

    if (test_lab = Cucumber::Chef::TestLab.new)
      if (provider = test_lab.create)
        if (provisioner = Cucumber::Chef::Provisioner.new(test_lab))

          provisioner.build

          puts
          puts("If you are using AWS, be sure to log into the chef-server webui and change the default admin password at least.")
          puts
          puts("Your test lab has now been provisioned!  Enjoy!")
          puts
          test_lab.status

        else
          puts(set_color("Could not create the provisioner!", :red, true))
        end
      else
        puts(set_color("Could not create the server!", :red, true))
      end
    else
      puts(set_color("Could not create a new instance of test lab!", :red, true))
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e.message)
  end

################################################################################
# DESTROY
################################################################################

  desc "destroy [container] [...]", "Destroy the cucumber-chef test lab or a single or multiple containers if specified"
  method_option :test, :type => :boolean, :desc => "INTERNAL USE ONLY"
  def destroy(*args)
    boot

    if (test_lab = Cucumber::Chef::TestLab.new) && test_lab.exists?
      if args.count == 0
        test_lab.status

        if yes?(set_color("Are you sure you want to destroy the test lab?", :red, true))
          puts
          puts(set_color("You have 5 seconds to abort!", :red, true))
          puts
          5.downto(1) do |x|
            print("#{x}...")
            sleep(1)
          end
          puts("BOOM!")
          puts

          ZTK::Benchmark.bench(:message => "Destroy #{Cucumber::Chef::Config.provider.upcase} instance '#{test_lab.id}'", :mark => "completed in %0.4f seconds.") do
            test_lab.destroy
          end
        else
          puts
          puts(set_color("Whew! That was close!", :green, true))
        end
      else
        if yes?(set_color("Are you sure you want to destroy the container#{args.count > 1 ? 's' : nil} #{args.collect{|a| "'#{a}'"}.join(', ')}?", :red, true))
          puts
          puts(set_color("You have 5 seconds to abort!", :red, true))
          puts
          5.downto(1) do |x|
            print("#{x}...")
            sleep(1)
          end
          puts("BOOM!")
          puts

          args.each do |container|
            ZTK::Benchmark.bench(:message => "Destroy container '#{container}'", :mark => "completed in %0.4f seconds.") do
              test_lab.containers.destroy(container)
            end
          end

        else
          puts
          puts(set_color("Whew! That was close!", :green, true))
        end

      end

    else
      raise Cucumber::Chef::Error, "We could not find a running test lab."
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

################################################################################
# UP
################################################################################

  desc "up", "Power up the cucumber-chef test lab"
  def up
    boot

    if (test_lab = Cucumber::Chef::TestLab.new) && test_lab.dead?
      ZTK::Benchmark.bench(:message => "Booting #{Cucumber::Chef::Config.provider.upcase} instance '#{test_lab.id}'", :mark => "completed in %0.4f seconds.") do
        test_lab.up
      end
    else
      raise Cucumber::Chef::Error, "We could not find a powered off test lab."
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

################################################################################
# DOWN
################################################################################

  desc "down", "Power off the cucumber-chef test lab"
  def down
    boot

    if (test_lab = Cucumber::Chef::TestLab.new) && test_lab.alive?
      ZTK::Benchmark.bench(:message => "Downing #{Cucumber::Chef::Config.provider.upcase} instance '#{test_lab.id}'", :mark => "completed in %0.4f seconds.") do
        test_lab.down
      end
    else
      raise Cucumber::Chef::Error, "We could not find a running test lab."
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

################################################################################
# RELOAD
################################################################################

  desc "reload", "Reload the cucumber-chef test lab"
  def reload
    boot

    if (test_lab = Cucumber::Chef::TestLab.new) && test_lab.alive?
      ZTK::Benchmark.bench(:message => "Reloading #{Cucumber::Chef::Config.provider.upcase} instance '#{test_lab.id}'", :mark => "completed in %0.4f seconds.") do
        test_lab.reload
      end
    else
      raise Cucumber::Chef::Error, "We could not find a running test lab."
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

  desc "genmac", "Generate an RFC compliant private MAC address"
  def genmac
    boot

    puts Cucumber::Chef::Containers.generate_mac
  end

  desc "genip", "Generate an RFC compliant private IP address"
  def genip
    boot

    puts Cucumber::Chef::Containers.generate_ip
  end

################################################################################
# STATUS
################################################################################

  desc "status", "Displays the current status of the test lab."
  method_option :containers, :type => :boolean, :desc => "Display container status.", :default => false
  method_option :attributes, :type => :boolean, :desc => "Display chef-client attributes for containers.", :default => false
  method_option :test, :type => :boolean, :desc => "INTERNAL USE ONLY"
  def status
    boot

    if (test_lab = Cucumber::Chef::TestLab.new)
      if @options.containers?
        if test_lab.alive?

          if test_lab.containers.count > 0
            headers = [:name, :alive, :distro, :ip, :mac, :"chef version", :persist]
            results = ZTK::Report.new.spreadsheet(Cucumber::Chef::Container.all, headers) do |container|
              chef_version = "N/A"
              alive = (test_lab.bootstrap_ssh(:ignore_exit_status => true).exec(%(ping -n -c 1 -W 1 #{container.ip}), :silence => true).exit_code == 0)
              if alive
                chef_version = test_lab.proxy_ssh(container.id, :ignore_exit_status => true).exec(%(/usr/bin/env chef-client -v), :silence => true).output.chomp
              end

              OpenStruct.new(
                :name => container.id,
                :ip => container.ip,
                :mac => container.mac,
                :distro => container.distro,
                :alive => alive,
                :"chef version" => chef_version,
                :persist => container.persist,
                :chef_attributes => container.chef_client
              )
            end

            if @options.attributes?
              results.rows.each do |result|
                puts
                puts("-" * results.width)
                puts("Chef-Client attributes for '#{result.name.to_s.downcase}':")
                puts("-" * results.width)
                puts(JSON.pretty_generate(result.chef_attributes))
              end
            end
          else
            raise Cucumber::Chef::Error, "We could not find any containers!"
          end

        else
          raise Cucumber::Chef::Error, "We could not find a running test lab."
        end

      else
        test_lab.status
      end
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e.message)
  end

################################################################################
# SSH
################################################################################

  desc "ssh [container]", "SSH to cucumber-chef test lab or [container] if specified"
  method_option :bootstrap, :type => :boolean, :desc => "Use the bootstrap settings", :default => false
  def ssh(*args)
    boot

    if (test_lab = Cucumber::Chef::TestLab.new) && test_lab.alive?
      if @options.bootstrap?
        puts([set_color("Attempting bootstrap SSH connection to cucumber-chef '", :blue, true), set_color("test lab", :cyan, true), set_color("'...", :blue, true)].join)
        test_lab.bootstrap_ssh.console
      elsif args.size == 0
        puts([set_color("Attempting SSH connection to the '", :blue, true), set_color("test lab", :cyan, true), set_color("'...", :blue, true)].join)
        test_lab.ssh.console
      elsif args.size > 0
        container = args[0]
        puts([set_color("Attempting proxy SSH connection to the container '", :blue, true), set_color(container, :cyan, true), set_color("'...", :blue, true)].join)
        test_lab.proxy_ssh(container).console
      else
        raise Cucumber::Chef::Error, "You did not specify a valid combination of options."
      end
    else
      raise Cucumber::Chef::Error, "We could not find a running test lab."
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

################################################################################
# PS
################################################################################

  desc "ps [ps-options]", "Snapshot of the current cucumber-chef test lab container processes."
  def ps(*args)
    boot

    if (test_lab = Cucumber::Chef::TestLab.new) && test_lab.alive?
      puts("-" * 80)
      test_lab.ssh.exec("lxc-ps --lxc -- #{args.join(" ")}")
      puts("-" * 80)
    else
      raise Cucumber::Chef::Error, "We could not find a running test lab."
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

################################################################################
# LOG
################################################################################

  desc "log", "Streams the cucumber-chef local and test lab logs to the terminal."
  def log
    boot

    if ($test_lab = Cucumber::Chef::TestLab.new) && $test_lab.exists? && $test_lab.alive?
      $tail_thread_remote = Thread.new do
        $test_lab.ssh.exec("tail -n 0 -f /home/#{$test_lab.ssh.config.user}/.cucumber-chef/cucumber-chef.log")
      end

      log_file = File.open(Cucumber::Chef.log_file, "r")
      log_file.seek(0, ::IO::SEEK_END)
      loop do
        if !(data = (log_file.readline rescue nil)).nil?
          print(data)
        else
          sleep(1)
        end
      end

      $tail_thread_remote.join
    else
      raise Cucumber::Chef::Error, "We could not find a running test lab."
    end

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

################################################################################
# DIAGNOSE
################################################################################

  desc "diagnose <container>", "Provide diagnostics from the chef-client on the specified container."
  method_option :strace, :type => :boolean, :desc => "output the chef-client 'chef-stacktrace.out'", :aliases => "-s", :default => true
  method_option :log, :type => :boolean, :desc => "output the chef-client 'chef.log'", :aliases => "-l", :default => true
  method_option :lines, :type => :numeric, :desc => "output the last N lines of the chef-client 'chef.log'", :aliases => "-n", :default => 1
  def diagnose(container)
    boot

    if (test_lab = Cucumber::Chef::TestLab.new) && test_lab.alive?
      puts([set_color("Attempting to collect diagnostic information on cucumber-chef container '", :blue, true), set_color(container, :cyan, true), set_color("'...", :blue, true)].join)
      if @options.strace?
        puts
        puts("chef-stacktrace.out:")
        puts(set_color("============================================================================", :bold))
        test_lab.proxy_ssh(container).exec("[[ -e /var/chef/cache/chef-stacktrace.out ]] && cat /var/chef/cache/chef-stacktrace.out")
        print("\n")
      end
      if @options.log?
        puts
        puts("chef.log:")
        puts(set_color("============================================================================", :bold))
        test_lab.proxy_ssh(container).exec("[[ -e /var/log/chef/client.log ]] && tail -n #{@options.lines} /var/log/chef/client.log")
      end
    else
      raise Cucumber::Chef::Error, "We could not find a running test lab."
    end
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

################################################################################
# DISPLAYCONFIG
################################################################################

  desc "displayconfig", "Display the current cucumber-chef config."
  method_option :test, :type => :boolean, :desc => "INTERNAL USE ONLY"
  def displayconfig
    boot

    details = {
      "root_dir" => Cucumber::Chef.root_dir,
      "home_dir" => Cucumber::Chef.home_dir,
      "log_file" => Cucumber::Chef.log_file,
      "artifacts_dir" => Cucumber::Chef.artifacts_dir,
      "config_rb" => Cucumber::Chef.config_rb,
      "labfile" => Cucumber::Chef.labfile,
      "chef_repo" => Cucumber::Chef.chef_repo,
      "chef_user" => Cucumber::Chef.chef_user,
      "chef_identity" => Cucumber::Chef.chef_identity,
      "bootstrap_user" => Cucumber::Chef.bootstrap_user,
      "bootstrap_user_home_dir" => Cucumber::Chef.bootstrap_user_home_dir,
      "bootstrap_identity" => Cucumber::Chef.bootstrap_identity,
      "lab_user" => Cucumber::Chef.lab_user,
      "lab_user_home_dir" => Cucumber::Chef.lab_user_home_dir,
      "lab_identity" => Cucumber::Chef.lab_identity,
      "lxc_user" => Cucumber::Chef.lxc_user,
      "lxc_user_home_dir" => Cucumber::Chef.lxc_user_home_dir,
      "lxc_identity" => Cucumber::Chef.lxc_identity,
      "chef_pre_11" => Cucumber::Chef.chef_pre_11
    }
    max_key_length = details.collect{ |k,v| k.to_s.length }.max

    puts("-" * 80)
    say(Cucumber::Chef::Config.configuration.to_yaml, :bold)
    puts("-" * 80)
    details.each do |key,value|
      puts("%#{max_key_length}s = %s" % [key.downcase, value.inspect])
    end
    puts("-" * 80)
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e.message)
  end

################################################################################
# CREATE
################################################################################

  desc "create <project>" , "Create a project template for testing an infrastructure."
  def create(project)

    create_project(project)
    root_dir = Cucumber::Chef.locate_parent(".chef")
    features_dir = File.join(root_dir, "features")
    feature = File.join(features_dir, "#{project}.feature")
    steps = File.join(features_dir, "step_definitions", "#{project}.steps")

    puts
    puts(set_color("Project created!", :green, true))
    say("Please look at the README in '#{features_dir}/#{project}/', and the example features (#{File.basename(feature)}) and steps (#{File.basename(steps)}), which I have autogenerated for you.", :green)
    puts

  rescue Cucumber::Chef::Error => e
    $logger.fatal { e.backtrace.join("\n") }
    fatal(e)
  end

################################################################################
# DEPRECIATED TASKS
################################################################################

  depreciated_tasks = {
    "teardown" => "You should execute the 'destroy' task instead.",
    "info" => "You should execute the 'status' task instead.",
    "test" => "You should execute 'cucumber' or 'rspec' directly."
  }

  depreciated_tasks.each do |old_task, message|
    desc old_task, "*DEPRECIATED* - #{message}"
    define_method(old_task) do
      puts
      puts(set_color("The '#{old_task}' task is *DEPRECIATED* - #{message}", :red, true))
      puts
    end
  end

################################################################################

end

CucumberChef.start