#!/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)

      get_aws_credentials

      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

    def get_aws_credentials
      say "OHAI!", :magenta
      puts
      say "Cucumber-Chef uses Amazon Web Services to build a test lab for automated infrastructure testing."
      say "I need a few details before I can set up your test lab for you."
      puts
      say "We're going to use symmetric keys to authenticate with the AWS API."
      say "First, I need your access key."
      say "Your access key identifies you as you make API calls.  It's not a secret."
      say "You can find it under 'Access Credentials', on https://aws-portal.amazon.com/gp/aws/securityCredentials"
      puts
      @aws_access_key = ask "What is your AWS Access Key?", :bold
      puts
      say "Now I need your secret access key.  This *is* a secret.  The clue's in the name."
      say "This is just a string of characters used to create the digital signature included in your API requests."
      say "You can also find this under 'Access Credentials', on https://aws-portal.amazon.com/gp/aws/securityCredentials"
      puts
      @aws_secret_access_key = ask "What is your AWS Secret Access Key?", :bold
      puts
      say "Right. Now I need to know about the ssh key pair you want to use to connect to EC2 machines."
      say "I need the name of the key pair. You can see this on the AWS management console, under Network & Security > Key Pairs"
      puts
      @aws_ssh_id = ask "What is your AWS Key Pair called?", :bold
      puts
      say "I also need to know what the ssh key is called - the actual name of the file on your local machine, eg #{@aws_ssh_id}.pem"
      puts
      @aws_ssh_key = ask "What's the filename of your ssh key?", :bold
      puts
      say "And, finally, I need to know where you keep it, on the file system. Often this is ~/.ssh"
      puts
      @aws_ssh_key_dir = ask "What directory contains your ssh key?", :bold
      puts
      say "OK, nearly there. AWS uses different keys depending on which region you use."
      say "For example, 'us-east', 'us-west', or 'eu-west'"
      puts
      @region = ask "Which region are you using?", :bold
      puts
      say("One last thing.  If you're using librarian-chef, we want to be sure all the hooks are in place.")
      puts
      @librarian_chef = yes?("Does this chef-repo use librarian-chef?", :bold)
      puts
      say "Awesome.  Thank you!"
      puts
    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,
      "bootstrap_identity" => Cucumber::Chef.bootstrap_identity,
      "chef_user" => Cucumber::Chef.chef_user,
      "chef_identity" => Cucumber::Chef.chef_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