#!/usr/bin/env ruby
# coding: utf-8
require 'rainbow'
require 'peach'
require 'json'
require 'clamp'
require 'net/ssh'
require 'net/scp'

def infohelper(nodes, parseable, grepable)
  if parseable
    puts nodes.to_json
  else
    nodes.each do |node|
      host = "#{node['role']}:#{node['environment']}:#{node['instance']}"
      unless grepable
        puts Rainbow(host).green
      end
      keys = node.keys.sort
      keys.each do |key|
        value = node[key]
        if grepable
          puts "#{host}|#{key}|#{value}"
        else
          value = Time.at(node[key].to_i) if key == 'launch_time'
          unless key == keys.last
            puts "  ┠  #{Rainbow(key).cyan}: #{value}"
          else
            puts "  ┖  #{Rainbow(key).cyan}: #{value}\n\n"
          end
        end
      end
    end
  end
end

def sshcmd(node, commands)
  Net::SSH.start(
    node['hostname'],
    'admin',
    :key_data => [$api.ssh(node['role'], node['environment'], node['instance'])['key']],
    :config => false,
    :keys_only => true,
    :paranoid => false
  ) do |ssh|
    exit_code = nil
    exit_signal = nil
    commands.each do |command|
      ssh.open_channel do |channel|
        channel.exec(command) do |chx, success|
          unless success
            abort "FAILED: couldn't execute command (ssh.channel.exec)"
          end
          channel.on_data do |chd,data|
            puts "#{Rainbow(node['role']).yellow}:#{Rainbow(node['environment']).yellow}:#{Rainbow(node['instance']).yellow}> #{data}"
          end

          channel.on_extended_data do |che,type,data|
            puts "#{Rainbow(node['role']).yellow}:#{Rainbow(node['environment']).yellow}:#{Rainbow(node['instance']).red}> #{data}"
          end

          channel.on_request("exit-status") do |chs,data|
            exit_code = data.read_long
            if exit_code != 0
              exit exit_code
            end
          end

          channel.on_request("exit-signal") do |chg, data|
            exit_signal = data.read_string
          end
        end
      end
    end
  end
end

def split_attrs(attribute_list)
  opts = {}
  attribute_list.each do |attr_|
    key, value = attr_.split('=', 2)
    split = key.split('.')
    cur = opts
    split.each_with_index do |part, idx|
      if idx == split.size - 1
        # leaf, add the value
        cur[part] = value
      else
        cur[part] ||= {}
      end
      cur = cur[part]
    end
  end
  opts
end

module Gaptool
  class InitCommand < Clamp::Command
    option ["-r", "--role"], "ROLE", "Resource name to initilize", :required => true
    option ["-e", "--environment"], "ENVIRONMENT", "Which environment, e.g. production", :required => true
    option ["-z", "--zone"], "ZONE", "AWS availability zone to put node in", :default => 'us-west-2c'
    option ["-t", "--type"], "TYPE", "Type of instance, e.g. m1.large", :required => true
    option ["-s", "--security-group"], "SECURITY_GROUP", "Security group name. Defaults to $role-$environment", required: false
    option ['-a', '--ami'], "AMI_ID", "Use a specific AMI for the instance (i.e. ami-xxxxxxx)", :required => false
    option ["-C", "--chef-repo"], "GITURL", "git url for the chef repository", :requireds => false
    option ["-b", "--chef-branch"], "BRANCH", "branch of the chef repository to use", :required => false
    option ['-R', "--chef_runlist"], 'RECIPE|ROLE', 'override chef run_list. recipe[cb::recipe] or role[myrole]. Can be specified multiple times', :multivalued => true, :attribute_name => 'chef_runlist'
    option ['--no-terminate'], :flag, "Protect this instance from being terminated"
    def execute
      no_terminate = no_terminate? ? true : nil
      $api.addnode(zone, type, role, environment, nil,
                   security_group, ami, chef_repo, chef_branch, chef_runlist,
                   no_terminate)
    end
  end

  class TerminateCommand < Clamp::Command
    option ["-i", "--instance"], "INSTANCE", "Instance ID, e.g. i-12345678", :required => true
    option ["-z", "--zone"], "ZONE", "AWS region of the node (deprecated/ignored)", :default => nil
    option ["-r", "--role"], "ROLE", "Resource name to initilize", :required => true
    option ["-e", "--environment"], "ENVIRONMENT", "Which environment, e.g. production", :required => true
    def execute
      node = $api.getonenode(instance)
      nodes = [node]
      infohelper(nodes, false, false)
      zone = node['zone'][0..-2]
      if node['environment'] != environment || node['role'] != role
        puts Rainbow("'#{node['role']}-#{node['environment']}' do not match provided value (#{role}-#{environment})").red
        exit 1
      end
      if node['terminate'] == 'false'
        puts Rainbow('"terminate" command is disabled for this instance').red
        exit 2
      end
      print Rainbow("Terminate instance? [type yes to confirm]: ").green
      res = $stdin.gets.chomp
      if res.downcase == 'yes'
        puts "Terminating instance #{node['role']}:#{node['environment']}:#{node['instance']} in region #{zone}"
        begin
          res = $api.terminatenode(instance, zone)
        rescue
          puts Rainbow('Cannot terminate instance').red
        end
      end
    end
  end

  class RuncmdCommand < Clamp::Command
    option ["-r", "--role"], "ROLE", "Role name to ssh to", :required => true
    option ["-e", "--environment"], "ENVIRONMENT", "Which environment, e.g. production", :required => true
    option ["-i", "--instance"], "INSTANCE", "Instance ID, e.g. i-12345678", :required => false
    parameter "COMMAND ...", "Command to run", :attribute_name => :command
    def execute
      if !instance.nil?
        node = $api.getonenode(instance)
        nodes = [$api.getonenode(instance)]
      else
        nodes = $api.getenvroles(role, environment)
      end
      nodes.peach do |node|
        commands = [ command.join(' ') ]
        sshcmd(node, commands)
      end
    end
  end

  class SshCommand < Clamp::Command
    option ["-r", "--role"], "ROLE", "Role name to ssh to", :required => true
    option ["-e", "--environment"], "ENVIRONMENT", "Which environment, e.g. production", :required => true
    option ["-i", "--instance"], "INSTANCE", "Node instance, leave blank to query avilable nodes", :require => false
    option ["-f", "--first"], :flag, "Just connect to first available instance"
    option ["-t", "--tmux"], :flag, "Open cluster in windows in a tmux session"

    def execute
      if tmux?
        nodes = $api.getenvroles(role, environment)
        system "tmux start-server"
        nodes.each_index do |i|
          @ssh = $api.ssh(role, environment, nodes[i]['instance'])
          if i == 0
            system "tmux new-session -d -s #{role}-#{environment} -n #{nodes[i]['instance']}"
          else
            system "tmux new-window -t #{role}-#{environment}:#{i} -n #{nodes[i]['instance']}"
          end
          File.open("/tmp/gtkey-#{nodes[i]['instance']}", 'w') {|f| f.write(@ssh['key'])}
          File.chmod(0600, "/tmp/gtkey-#{nodes[i]['instance']}")
          system "tmux send-keys -t #{role}-#{environment}:#{i} 'SSH_AUTH_SOCK=\"\" ssh -i /tmp/gtkey-#{nodes[i]['instance']} admin@#{@ssh['hostname']}' C-m"
        end
        system "tmux attach -t #{role}-#{environment}"
      else
        if instance
          @ssh = $api.ssh(role, environment, instance)
        else
          nodes = $api.getenvroles(role, environment)
          if first? || nodes.size == 1
            puts "No instance specified, but only one instance in cluster or first forced"
            @ssh = $api.ssh(role, environment, nodes.first['instance'])
          else
            puts "No instance specified, querying list."
            nodes.each_index do |i|
              puts "#{i}: #{nodes[i]['instance']}"
            end
            print Rainbow("Select a node: ").cyan
            @ssh = $api.ssh(role, environment, nodes[$stdin.gets.chomp.to_i]['instance'])
          end
        end
        File.open('/tmp/gtkey', 'w') {|f| f.write(@ssh['key'])}
        File.chmod(0600, '/tmp/gtkey')
        system "SSH_AUTH_SOCK='' ssh -i /tmp/gtkey admin@#{@ssh['hostname']}"
      end
    end

  end

  class InfoCommand < Clamp::Command
    option ["-r", "--role"], "ROLE", "Role name, e.g. frontend", :required => false
    option ["-e", "--environment"], "ENVIRONMENT", "Which environment, e.g. production", :required => false
    option ["-i", "--instance"], "INSTANCE", "Node instance, leave blank to query avilable nodes", :required => false
    option ["-p", "--parseable"], :flag, "Display in non-pretty parseable JSON"
    option ["-g", "--grepable"], :flag, "Display in non-pretty grep-friendly text"
    option ['-H', '--hidden'], :flag, 'Display hidden hosts'

    def execute
      @nodes = Array.new
      params = hidden? ? {hidden: true} : {}
      if instance
        @nodes = [$api.getonenode(instance)]
      elsif role && environment
        @nodes = $api.getenvroles(role, environment, params)
      elsif role && !environment
        @nodes = $api.getrolenodes(role, params)
      else
        @nodes = $api.getallnodes(params)
      end
      infohelper(@nodes, parseable?, grepable?)
    end
  end

  class SetCommand < Clamp::Command
    option ["-i", "--instance"], "INSTANCE", "Node instance, required", required: true
    option ["-p", "--parseable"], :flag, "Display in non-pretty parseable JSON"
    option ["-g", "--grepable"], :flag, "Display in non-pretty grep-friendly text"
    option ['-k', '--parameter'], 'NAME', 'Set parameter for the node', required: true, multivalued: true
    option ['-v', '--value'], 'VALUE', 'Value for parameter', required: true, multivalued: true

    def convert_bool(v)
      case v.downcase
      when "true"
        true
      when "false"
        false
      else
        v
      end
    end

    def execute
      if parameter_list.length != value_list.length
        puts Rainbow("parameter and value length mismatch").red
      end
      params = Hash[parameter_list.each_with_index.map {|p, i| [p, convert_bool(value_list[i])]}]
      infohelper([$api.setparameters(instance, params)], parseable?, grepable?)
    end
  end

  class ChefrunCommand < Clamp::Command
    option ["-r", "--role"], "ROLE", "Role name to ssh to", :required => true
    option ["-e", "--environment"], "ENVIRONMENT", "Which environment, e.g. production", :required => true
    option ["-i", "--instance"], "INSTANCE", "Instance ID, e.g. i-12345678", :required => false
    option ["-A", "--attribute"], "ATTRIBUTE", "Pass one or more parameters to the deploy recipe in recipe.attr=value format", :multivalued => true
    option ["-b", "--chef-branch"], "BRANCH", "branch of the chef repository to use (defaults to last branch used during init/chefrun)", :default => nil

    def execute
      attrs = split_attrs(attribute_list)
      if !instance.nil?
        nodes = [$api.getonenode(instance)]
      else
        nodes = $api.getenvroles(role, environment)
      end
      nodes.peach do |node|
        if node['chef_runlist'].nil?
          runlist = ['recipe[main]']
        elsif node['chef_runlist'].is_a? Array
          runlist = node['chef_runlist']
        else
          runlist = eval(node['chef_runlist'])
        end
        json = {
          'this_server' => "#{role}-#{environment}-#{node['instance']}",
          'role' => role,
          'environment' => environment,
          'app_user' => node['appuser'],
          'run_list' => runlist,
          'hostname' => node['hostname'],
          'instance' => node['instance'],
          'zone' => node['zone'],
          'itype' => node['itype'],
          'apps' => eval(node['apps'] || '[]'),
        }.merge(attrs)
        git="git fetch --all; git reset --hard origin/`git rev-parse --abbrev-ref HEAD`"

        unless chef_branch.nil?
          json['chefbranch'] = chef_branch
          git = "git checkout -f #{chef_branch}; git fetch --all; git reset --hard origin/#{chef_branch}"
        end
        commands = [
          "cd ~admin/ops; #{git} ",
          "echo '#{json.to_json}' > ~admin/solo.json",
          "sudo chef-solo -c ~admin/ops/cookbooks/solo.rb -j ~admin/solo.json -E #{environment}"
        ]
        sshcmd(node, commands)
      end
    end
  end

  class DeployCommand < Clamp::Command
    option ["-a", "--app"], "APP", "Application to deploy", :required => true
    option ["-m", "--migrate"], :flag, "Toggle running migrations"
    option ["-e", "--environment"], "ENVIRONMENT", "Which environment, e.g. production", :required => true
    option ["-b", "--branch"], "BRANCH", "Git branch to deploy, default is master", :required => false
    option ["-r", "--rollback"], :flag, "Toggle this to rollback last deploy"
    option ["-i", "--instance"], "INSTANCE", "Instance ID, e.g. i-12345678", :required => false
    option ["-A", "--attribute"], "ATTRIBUTE", "Pass one or more parameters to the deploy recipe in recipe.attr=value format", :multivalued => true
    option ['-H', '--hidden'], :flag, 'Display hidden hosts'

    def execute
      attrs = split_attrs(attribute_list)
      if instance
        n = $api.getonenode(instance)
        if n['environment'] != environment
          abort "Instance #{instance} is not in environment #{environment}"
        elsif !n['apps'].include?(app)
          abort "Instance #{instance} does not host #{app} in env #{environment}"
        end
        nodes = [n]
      else
        params = hidden? ? {hidden: true} : {}
        nodes = $api.getappnodes(app, environment, params)
      end
      nodes.peach do |node|
        if node['chef_runlist'].nil?
          runlist = ['recipe[deploy]']
        elsif node['chef_runlist'].is_a? Array
          runlist = node['chef_runlist']
        else
          runlist = eval(node['chef_runlist'])
        end
        json = {
          'this_server' => "#{node['role']}-#{environment}-#{node['instance']}",
          'role' => node['role'],
          'environment' => environment,
          'app_user' => node['appuser'],
          'run_list' => runlist,
          'hostname' => node['hostname'],
          'instance' => node['instance'],
          'zone' => node['zone'],
          'itype' => node['itype'],
          'apps' => eval(node['apps']),
          'app_name' => app,
          'app' => app,
          'rollback' => rollback?,
          'branch' => branch || 'master',
          'migrate' => migrate?
        }.merge(attrs).to_json
        commands = [
          "cd ~admin/ops; git pull",
          "echo '#{json}' > ~admin/solo.json",
          "sudo chef-solo -c ~admin/ops/cookbooks/solo.rb -j ~admin/solo.json -E #{environment}"
        ]
        sshcmd(node, commands)
      end
    end
  end

  class RehashCommand < Clamp::Command
    option ["-y", "--yes"], :flag, "YES I REALLY WANT TO DO THIS"
    def execute
      if yes?
        puts $api.rehash()
      else
        puts "You need to run this with -y\nIf you don't know what this does or aren't sure, DO NOT RUN IT\nThis will regenerate all host metadata on gaptool-server\nand can break in-progress operations."
      end
    end
  end

  class VersionCommand < Clamp::Command
    option ["-r", "--remote"], :flag, "Include remote API version"
    def execute
      version = File.read(File.realpath(File.join(File.dirname(__FILE__), "..", 'VERSION'))).strip
      puts "gaptool-client #{version} using gaptool-api #{$api.version}"
      if remote?
        vinfo = $api.api_version()
        puts "gaptool-server #{vinfo['server_version']}"
      end
    end
  end

  class MainCommand < Clamp::Command

    subcommand "info", "Displays information about nodes", InfoCommand
    subcommand "init", "Create new application cluster", InitCommand
    subcommand "terminate", "Terminate instance", TerminateCommand
    subcommand "ssh", "ssh to cluster host", SshCommand
    subcommand "set", "update properties for a node", SetCommand
    subcommand "chefrun", "chefrun on a resource pool", ChefrunCommand
    subcommand "deploy", "deploy on an application", DeployCommand
    subcommand "rehash", "Regenerate all host metadata. KNOW WHAT THIS DOES BEFORE RUNNING IT", RehashCommand
    subcommand "runcmd", "Run command on instance", RuncmdCommand
    subcommand "version", "Show version", VersionCommand

  end
end

Gaptool::MainCommand.run