#!/usr/bin/env ruby

$:.unshift File.expand_path("../../lib", __FILE__)

require 'stacco'
require 'pathname'
require 'highline'
require 'inifile'
require 'digest'

class String
  def colored(color)
    color_number = (color.kind_of?(Integer) ? color : COLORS[color]) + 30
    "\x1b[#{color_number}m#{self}\x1b[0m"
  end

  def guess_color
    return :black if self.nil?
    return :cyan unless self.kind_of? String

    case self
    when /DOWN$/
      :red
    when /COMPLETE$/
      :green
    when /IN_PROGRESS$/
      :yellow
    when /WORKING$/
      :yellow
    when /CHANGED$/
      :blue
    when /started$/
      :blue
    when /FAILED$/
      :red
    when /^CREATE/
      :green
    when /^DELETE/
      :blue
    when /^ROLLBACK/
      :red
    else
      :white
    end
  end

  COLORS = {
    black: 0,
    red: 1,
    green: 2,
    yellow: 3,
    blue: 4,
    magenta: 5,
    cyan: 6,
    white: 7
  }
end

def watch_events_on(stack, opts = {})
  terminal_statuses = [opts.delete(:until)].flatten.compact
  archival_mode = opts.delete(:show_previous)

  return unless terminal_statuses.empty? or stack.operation_in_progress? or archival_mode

  _, terminal_width = `stty size`.chomp.split(" ").map{ |d| d.to_i }

  operation_start_time = stack.up_since
  #watch_began = Time.now

  trap('INT') do
    $stderr.puts
    Kernel.exit 0
  end

  begin
    stack.stream_events.each do |ev|
      status = ev.status
      color = status.guess_color

      if ev.resource_type == "AWS::CloudFormation::Stack"
        operation_start_time = ev.timestamp
        next unless ev.live or archival_mode

        status_part = "#{ev.operation} #{ev.logical_resource_id} #{status} at #{ev.timestamp}"
        puts
        puts "=== #{status_part.colored(color)} ".ljust(terminal_width, '=')
        puts

        break if ev.live and terminal_statuses.find{ |s| ev.resource_status =~ s }
      else
        next unless ev.live or (ev.op and ev.op.live) or archival_mode
        next if status == "WORKING" and not ev.error and not $verbose

        relative_timestamp = ev.timestamp - operation_start_time

        timestamp_part = "[#{("+%.1f" % relative_timestamp).rjust(12, ' ')}]"
        status_part = "[#{status}]".rjust(10, ' ')
        resource_id_part = "#{ev.logical_resource_id}".rjust(30, ' ')
        action_part = "#{status_part} #{resource_id_part}"

        puts "#{timestamp_part} #{action_part.colored(color)} #{ev.error}"

        if ev.detail
          indent = " " * ((action_part.length + timestamp_part.length)/2)

          puts
          puts
          ev.detail.each do |detail_ln|
            puts "#{indent}#{detail_ln}"
          end
          puts
        end
      end

      $stdout.flush
    end
  rescue AWS::CloudFormation::Errors::ValidationError
    puts "Stack terminated."
    false
  end
end

def watch_operation_on(stack)
  return if $async

  if $watch
    watch_events_on(stack, until: /_COMPLETE$/)
  else
    Kernel.sleep 2 while stack.operation_in_progress?
  end
end

$verbose = true if ARGV.delete('-v')
$force = true if ARGV.delete('-f')
$watch = true if (ARGV.delete('-w') or $stdout.tty?)
$async = true if ARGV.delete('-a')

if stack_name_index = ARGV.index('-s')
  $stack_name = ARGV[stack_name_index + 1]
  ARGV[stack_name_index, 2] = []
end

unless ARGV.length >= 1
  $stderr.puts "usage: #{File.basename($0)} <subcommand ...>"
  Kernel.exit 1
end

console = HighLine.new

subcommand = ARGV.shift.intern

orch_config = {
  local: Pathname.new('orchestrator.yml'),
  user: Pathname.new(ENV['HOME']) + '.config' + 'stacco' + 'orchestrator.yml'
}
orch_config = orch_config[:local].file? ? orch_config[:local] : orch_config[:user]
orch_config.parent.mkpath

if subcommand == :init
  if $stdin.tty?
    new_orch_config = {aws: {}}

    aws_config = Pathname.new(ENV['HOME']) + '.aws' + 'config'
    if aws_config.file?
      aws_config = IniFile.new(filename: aws_config.to_s)
      aws_profiles = aws_config.sections - ["preview"]

      if aws_profiles.length > 1
        aws_profile = console.choose(aws_profiles) do |menu|
          menu.prompt = "Local AWS profile to source: "
        end
      else
        console.say "Sourcing AWS config from local AWS profile."
        aws_profile = aws_profiles.first
      end

      new_orch_config[:aws][:access_key_id] = aws_config[aws_profile]['aws_access_key_id']
      new_orch_config[:aws][:secret_access_key] = aws_config[aws_profile]['aws_secret_access_key']
      new_orch_config[:aws][:region] = aws_config[aws_profile]['region']
    end

    new_orch_config[:aws][:access_key_id] = console.ask("AWS account ID: ") do |question|
      question.default = new_orch_config[:aws][:access_key_id]
      question.validate = /^[0-9A-Z]{20}$/
    end

    new_orch_config[:aws][:secret_access_key] = console.ask("AWS account secret: ") do |question|
      question.default = new_orch_config[:aws][:secret_access_key]
      question.validate = /^[=+\/0-9A-Za-z]{32,64}$/
    end

    new_orch_config[:aws][:region] = console.choose(*(AWS.regions.map{ |r| r.name }.sort)) do |menu|
      menu.prompt = "Orchestrator storage-bucket region: "
      menu.default = new_orch_config[:aws][:region]
      menu.layout = :one_line
      menu.select_by = :name_or_index
    end

    new_orch_config[:storage_bucket] = console.ask("Orchestrator storage-bucket name: ") do |question|
      question.validate = /^[A-Za-z][-.0-9A-Za-z]+[A-Za-z0-9]$/
    end

    # convert symbol-keys to string-keys
    new_orch_config = JSON.parse(new_orch_config.to_json)
  else
    new_orch_config = YAML.load($stdin.read)
  end

  orch_config.open('w'){ |f| f.write(new_orch_config.to_yaml) }

  orch = Stacco::Orchestrator.from_config(orch_config)
  puts "orchestrator initialized"
  Kernel.exit 0
end

unless orch_config.file?
  $stderr.puts "orchestrator not initialized!"
  $stderr.puts "Did you run 'stacco init'?"
  Kernel.exit 1
end

orch = Stacco::Orchestrator.from_config(orch_config)
stacks = orch.stacks

case subcommand
when :"stacks:list"
  puts "known stacks:"
  stacks.each do |st|
    puts "  #{st.name.colored(st.status.guess_color)}"
  end
  Kernel.exit 0
when :"stacks:import"
  stack_defn = $stdin.read
  stack = orch.define_stack stack_defn
  puts "stack '#{stack.name}' imported"
  Kernel.exit 0
end

if $stack_name
  stack = stacks.find{ |st| st.name == $stack_name }

  unless stack
    $stderr.puts "unknown stack '#{stack_name}'"
    Kernel.exit 1
  end
else
  case stacks.length
  when 0
    $stderr.puts "no stacks defined!"
    $stderr.puts "run 'cat ./stack.yml | stacco stacks:import'"
    Kernel.exit 1
  when 1
    stack = stacks.first
  else
    $stderr.puts "ambiguous target stack"
    $stderr.puts "try '#{File.basename($0)} -s <stackname> #{subcommand}'"
    Kernel.exit 1
  end
end

case subcommand
when :up
  if ARGV.empty?
    $stderr.puts "usage: #{File.basename($0)} up <layers> | all"
    Kernel.exit 1
  end

  stack.cancel_operation if $force

  stack.enable_layers( ARGV.delete('all') ? stack.available_layers : ARGV )
  unless stack.up!
    $stderr.puts "No changes."
    Kernel.exit 1
  end
  watch_operation_on stack

when :down
  if ARGV.empty?
    $stderr.puts "usage: #{File.basename($0)} down <layers> | all"
    Kernel.exit 1
  end

  stack.cancel_operation if $force
  watch_operation_on stack

  stack.disable_layers( ARGV.delete('all') ? stack.available_layers : ARGV )

  if stack.enabled_layers.empty?
    stack.down!
  else
    unless stack.up!
      $stderr.puts "No changes."
      Kernel.exit 1
    end
  end
  watch_operation_on stack

when :bounce
  if ARGV.empty?
    $stderr.puts "usage: #{File.basename($0)} bounce <layers> | all"
    Kernel.exit 1
  end

  stack.cancel_operation if $force
  watch_operation_on stack

  layers_for_op = ARGV.delete('all') ? stack.available_layers : ARGV

  stack.disable_layers layers_for_op
  stack.up!
  watch_operation_on stack

  stack.enable_layers layers_for_op
  stack.up!
  watch_operation_on stack

when :retry
  stack.cancel_operation if $force
  watch_operation_on stack

  stack.up!
  watch_operation_on stack

when :cancel
  unless stack.operation_in_progress?
    $stderr.puts "no operations in progress"
    Kernel.exit 1
  end

  stack.cancel_operation
  watch_operation_on stack

when :export
  puts stack.config.to_yaml

when :edit
  editor = ENV['EDITOR'] || 'vi'
  tmp_stackdef_path = Pathname.new("/tmp/stacco-stack.#{stack.name}.#{rand(32 ** 16).to_s(32)}")
  tmp_stackdef_path.open('w'){ |f| f.write(stack.config.to_yaml) }
  Kernel.system(editor, tmp_stackdef_path.to_s)
  stack.config = YAML.load(tmp_stackdef_path.read)


when :"template:body"
  require 'pp'

  stack.bake_template do |body, params|
    puts body
    false
  end

when :"config:list", :"config:show", :config
  require 'pp'

  terminal_height, terminal_width = `stty size`.chomp.split(" ").map{ |d| d.to_i }
  width_with_margin = terminal_width - 10

  stack.bake_template do |body, params|
    params.each do |k,v|
      if v.length >= 40
        v = v.gsub(/\s+/, ' ')
        v_digest = Digest::SHA1.digest(v).unpack('I')[0].to_s(32)
        v_desc = "#{v.length.to_s.colored(:green)} bytes = #{v_digest.colored(:blue)}"
        v_partial = "#{v[0...width_with_margin].colored(:yellow)}..."
        puts "#{k}:"
        puts "  #{v_desc}"
        puts "  #{v_partial}"
        puts
      else
        puts "#{k}: #{v.to_json.colored(:red)}"
      end
    end
    false
  end

when :"config:get"
  param_key = ARGV.shift

  stack.bake_template do |body, params|
    unless params.has_key?(param_key)
      $stderr.puts "Unset stack parameter '#{param_key}'".colored(:red)
      Kernel.exit 1
    end

    puts params[param_key]
    false
  end


when :"template:validate"
  success, err, pos = stack.validate
  if success
    puts "template is valid"
  else
    $stderr.puts err.colored(:red)
    if pos
      $stderr.puts
      line, column = pos
      line_before, line_after = line[0...column], line[column..-1]
      if line_after
        token_at, rest_of_line = line_after.split(' ', 2)
        $stderr.puts line_before + (token_at || '').colored(:red) + (rest_of_line || '')
      else
        $stderr.puts line_before + ' #'.colored(:blue) + ' <-'.colored(:red)
      end
      $stderr.puts
    end
  end

when :repl
  require 'irb'
  AWS.config(stack.aws_credentials)
  Stack = stack
  IRB.start

when :"cloudfront:init"
  stack.initialize_distributions!

when :"cloudfront:invalidate"
  unless dist_name = ARGV.shift and ARGV.length.nonzero?
    die "usage: #{File.basename($0)} cloudfront:invalidate <distribution> [paths]"
  end

  stack.invalidate_distributed_objects!(dist_name, ARGV)

when :connect
  conns = stack.connections

  bastion_host_name, bastion_host = conns.find{|k,v| k =~ /bastion/i and v.dns_name }

  unless host_vague_name = ARGV.shift
    if bastion_host
      conns.each do |resource_name, inst|
        connection_path = inst.dns_name ? [inst.dns_name] : [bastion_host.dns_name, inst.private_dns_name]
        puts "#{resource_name}: #{connection_path.join(" -> ".colored(:yellow))}"
      end
    else
      conns.find_all{ |rn, inst| inst.dns_name }.each do |(resource_name, inst)|
        puts "#{resource_name}: #{inst.dns_name}"
      end
    end

    Kernel.exit 0
  end

  if host_vague_name.include? '.'
    host_type = :physical
    host_physical_name = host_vague_name
  else
    host_type = :logical
    host_logical_name = host_vague_name
  end

  if host_type == :logical
    unless conns.has_key? host_logical_name
      $stderr.puts "#{host_logical_name}: not available"
      Kernel.exit 1
    end

    connect_to_host = conns[host_logical_name]
  else
    connect_to_host = OpenStruct.new(
      dns_name: nil,
      private_dns_name: host_physical_name,
      private_ip_address: host_physical_name
    )
  end

  possible_paths = [
    [connect_to_host],
    [bastion_host, connect_to_host]
  ]

  extant_paths = possible_paths.find_all(&:all?)

  extant_routes = extant_paths.map do |(first_host, *rest_of_hosts)|
    [first_host.public_dns_name] + rest_of_hosts.map(&:private_ip_address)
  end

  routable_path = extant_routes.find{ |f| f.compact.length.nonzero? }

  ssh_stanzas = routable_path.map do |hostname|
    ['ssh', '-A', '-tt', '-q',
      '-o', 'StrictHostKeyChecking=no',
      '-o', 'UserKnownHostsFile=/dev/null',
      "ubuntu@#{hostname}"]
  end

  ssh_cmd_parts = ssh_stanzas.flatten + ARGV

  Kernel.system 'ssh-add', stack.iam_private_key.local_path.to_s, out: '/dev/null', err: '/dev/null'
  Kernel.exec *ssh_cmd_parts

when :"status:wait"
  desired_status = ARGV.shift
  if desired_status == 'down'
    Kernel.sleep(2) while stack.up?
  else
    Kernel.sleep(2) until stack.up? and stack.aws_status == desired_status
  end

when :status
  stack.must_be_up!
  now = Time.now

  showing_task_layers = ARGV.delete('-a')

  puts "#{stack.description.colored(stack.aws_status.guess_color)} (state: #{stack.aws_status})"
  puts

  visible_layers = stack.available_layers
  (visible_layers = visible_layers.delete_if{ |l| l.task? }) unless showing_task_layers

  visible_layers.sort_by{ |layer| layer.full_name }.each do |layer|
    colored_resource_names = layer.resource_summaries.sort_by{ |res| res[:logical_resource_id] }.map{ |res| res[:logical_resource_id].colored(res[:resource_status].guess_color) }
    case layer.state
    when :healthy
      puts "  #{layer.full_name.colored(:white)}: #{colored_resource_names.join(' ')}"
    when :zombie
      puts "  #{layer.full_name.colored(:red)}: #{colored_resource_names.join(' ')}"
    when :queued
      puts "  #{layer.full_name.colored(:yellow)}"
    when :down
      puts "  #{layer.full_name}"
    end
  end

when :"status:normalize"
  stack.must_be_up!
  now = Time.now

  layers_to_disable = []
  layers_to_enable = []

  stack.available_layers.sort_by{ |layer| layer.full_name }.each do |layer|
    case layer.state
    when :zombie
      layers_to_enable.push layer.name
    when :queued
      layers_to_disable.push layer.name
    end
  end

  stack.enable_layers(layers_to_enable) unless layers_to_enable.empty?
  stack.disable_layers(layers_to_disable) unless layers_to_disable.empty?

when :events
  stack.must_be_up!

  watch_events_on stack, show_previous: $verbose

  $stderr.puts
else
  $stderr.puts "unknown subcommand"
  Kernel.exit 1
end