#!/usr/bin/env ruby
require 'json'
require 'yaml'
require 'pp'
require 'open3'
Thread.abort_on_exception = true

def yellow(txt); "\e[0;33m#{txt}\e[0m"; end
def red(txt);    "\e[0;31m#{txt}\e[0m"; end
def green(txt);  "\e[0;32m#{txt}\e[0m"; end
def blue(txt);   "\e[0;34m#{txt}\e[0m"; end

def do_it(cmd)
  puts "[exec] #{cmd}"
  system(cmd)
  unless $?.success?
    puts red("FATAL :(")
    exit 
  end
end

def just_cmd(cmd, skip_exit_status = false)
  puts "[exec] #{cmd}"
  output = `#{cmd}`
  if $?.success? || skip_exit_status
    output.strip
  else
    puts red(output)
    puts red("FATAL :(")
    exit
  end
end

if ARGV.count < 2
  puts "Just deploy your apps with docker and consul. Nothing else."
  puts ""
  puts "  #{$0} <yml> <command>"
  puts ""
  puts "  deploy [hash]      - deploy new version of app"
  puts "  status             - status of app"
  puts "  stop               - stop app"
  exit
end

CONFIG = YAML.load(File.read(ARGV[0]))
COMMAND = ARGV[1]

MIN_PORT = 10_000
MAX_PORT = 50_000
Node = Struct.new(:name, :ip, :roles)
NODES = JSON.load(just_cmd("curl -s #{CONFIG["consul"]}/v1/catalog/nodes")).sort_by { |n| 
  n["Node"]
}.map { |n|
  Node.new(n["Node"], n["Address"], (n["Meta"] || {})["roles"].to_s.split(",").reject{ |s| s.to_s.empty? })
}

def nodes(constraint)
  constraint ||= {}
  out = NODES
  if constraint["role"]
    out = out.select { |n| n.roles.index(constraint["role"]) }
  end
  if constraint["name"]
    out = out.select { |n| n.name == constraint["name"] }
  end
  out
end

def node(constraint)
  nodes(constraint).sample
end


def deploy(target = nil)
  build_node = node(CONFIG["image"]["constraint"])
  puts "build_node=#{build_node.to_h.inspect}"

  puts blue("+++ CLONE or UPDATE repository")
  do_it "ssh #{CONFIG["user"]}@#{build_node.ip} bash <<EOF
    set -e
    tmpfile=\\$(mktemp /tmp/git-ssh-#{CONFIG["app"]}.XXXXXX)
    echo '#!/bin/sh' > \\$tmpfile
    echo 'exec /usr/bin/ssh -o StrictHostKeyChecking=no -i #{CONFIG["image"]["key"]} \"\\$@\"' >> \\$tmpfile
    chmod +x \\$tmpfile
    export GIT_SSH=\\$tmpfile
    if [ -d #{CONFIG["app"]}-cache ]; then
      echo update cache...
      cd #{CONFIG["app"]}-cache
      git checkout . && git clean -dfx && git checkout master && git pull
      git branch | grep -v master | xargs -r git branch -D
    else
      echo clone...
      git clone git@github.com:#{CONFIG["image"]["repo"]} #{CONFIG["app"]}-cache && cd #{CONFIG["app"]}-cache
    fi
    git checkout #{target || CONFIG["image"]["target"]}
    rm \\$tmpfile\nEOF"

  puts blue("+++ calculate HASH and VERSION")
  hash = ARGV[2] || `ssh #{CONFIG["user"]}@#{build_node.ip} 'cd #{CONFIG["app"]}-cache && git rev-parse HEAD'`.strip
  puts "hash: #{hash}"

  o = JSON.load(`curl -s https://#{CONFIG["registry"]}/v2/#{CONFIG["app"]}/tags/list`)
  tags = o.is_a?(Hash) && o["errors"] ? [] : o["tags"]
  puts "tags: #{JSON.dump(tags)}"

  unless tags.index(hash)
    puts blue("+++ BUILD image")
    do_it "ssh #{CONFIG["user"]}@#{build_node.ip} bash <<EOF
      set -e
      cd #{CONFIG["app"]}-cache
      #{(CONFIG["image"]["pre_build"] || []).join("\n")}
      docker build -t #{CONFIG["registry"]}/#{CONFIG["app"]}:#{hash} .
      docker push #{CONFIG["registry"]}/#{CONFIG["app"]}:#{hash}
      echo Done.\nEOF"
  end
  vv = Time.now.to_i
  new_service = "#{CONFIG["app"]}-#{vv}"

  (CONFIG["deploy"]["pre"] || []).each_with_index do |section, index|
    puts blue("+++ PRE: #{section.inspect}")
    n = node(section["constraint"] || CONFIG["deploy"]["constraint"])
    puts "node=#{n.inspect}"
    do_it "ssh #{CONFIG["user"]}@#{n.ip} docker run " +
      "--label app=#{CONFIG["app"]} " +
      "--label service=#{new_service} " +
      "--name #{CONFIG["app"]}-pre-#{vv}-#{index} " +
      "#{(CONFIG["env"] || {}).map { |k, v| "-e #{k}='#{v}'" }.join(" ")} " +
      "#{CONFIG["registry"]}/#{CONFIG["app"]}:#{hash} #{section["cmd"]}"
    do_it "ssh #{CONFIG["user"]}@#{n.ip} docker rm #{CONFIG["app"]}-pre-#{vv}-#{index}"
  end

  deploy_nodes = nodes(CONFIG["deploy"]["constraint"])

  puts blue("+++ PULL")
  pull_threads = []
  deploy_nodes.each do |node|
    pull_threads << Thread.new do
      host = "#{CONFIG["user"]}@#{node.ip}"
      Open3.popen2e("ssh #{host} docker pull #{CONFIG["registry"]}/#{CONFIG["app"]}:#{hash}") { |i,o,w|
        while line = o.gets do 
          puts "[#{node.name}] " + line
        end
        raise "FATAL" unless w.value.success?
      }
    end
  end
  pull_threads.each(&:join)

  plan = {}
  if CONFIG["deploy"]["per"]
    amount = CONFIG["deploy"]["per"]*deploy_nodes.count
  else
    amount = CONFIG["deploy"]["amount"]
  end

  amount.times do |index|
    ip = deploy_nodes[index % deploy_nodes.length].ip
    plan[ip] ||= []
    plan[ip].push(index)
  end

  puts blue("+++ DEPLOY")
  deploy_threads = []
  deploy_nodes.each do |node|
    deploy_threads << Thread.new do
      host = "#{CONFIG["user"]}@#{node.ip}"

      puts "[#{node.name}] run"
      (plan[node.ip] || []).each do |index|
        port = nil
        10.times do
          a = MIN_PORT + rand(MAX_PORT - MIN_PORT)
          if just_cmd("ssh #{host} ss -ln src :#{a} | fgrep -c ':#{a}'", true) == "0"
            port = a
            break
          end
        end
        raise "Dont find free port :-((" unless port 
        puts "[#{node.name}] port=#{port}"

        do_it("ssh #{host} docker run -d " + 
              "--label app=#{CONFIG["app"]} " +
              "--label service=#{new_service} " +
              "--name=#{CONFIG["app"]}-#{vv}-#{index} " +
              (CONFIG["env"] || {}).map { |k, v| "-e #{k}='#{v}'" }.join(" ") + " " +
              "--restart unless-stopped " +
              "-p #{node.ip}:#{port}:#{CONFIG["deploy"]["port"]} " +
              "#{CONFIG["registry"]}/#{CONFIG["app"]}:#{hash} " +
              "#{CONFIG["deploy"]["cmd"]}")

        do_it %(curl -s -X PUT #{node.ip}:8500/v1/agent/service/register -d'{
          "Id": "#{CONFIG["app"]}-#{vv}-#{index}",
          "Name": "#{new_service}",
          "Port": #{port},
          "Check": {
            "DeregisterCriticalServiceAfter": "20m",
            "Interval": "7s",
            "HTTP": "http://#{node.ip}:#{port}#{CONFIG["deploy"]["check_url"] || "/"}"
          }
        }')
      end

      puts blue("+++ CONSUL wait for #{node.name}")
      loop do
        statuses = JSON.load(just_cmd("curl -s #{node.ip}:8500/v1/health/service/#{CONFIG["app"]}-#{vv}")).select { |s|
          s["Node"]["Address"] == node.ip
        }.flat_map { |s| s["Checks"] }.map { |c| c["Status"] }
        puts "#{node.name} => #{statuses.inspect}"
        break if statuses.uniq == ["passing"] or (plan[node.ip] || []).empty?
        sleep 3
      end
    end
  end
  deploy_threads.each(&:join)

  current_service = just_cmd("curl #{CONFIG["consul"]}/v1/kv/apps/#{CONFIG["app"]}/service?raw")
  current_hash = just_cmd("curl #{CONFIG["consul"]}/v1/kv/apps/#{CONFIG["app"]}/hash?raw")
  puts "CURRENT_SERVICE=#{current_service}"
  puts "CURRENT_HASH=#{current_hash}"
  puts "NEW_SERVICE=#{new_service}"
  do_it "curl -X PUT #{CONFIG["consul"]}/v1/kv/apps/#{CONFIG["app"]}/service -d#{new_service}"
  do_it "curl -X PUT #{CONFIG["consul"]}/v1/kv/apps/#{CONFIG["app"]}/hash -d#{hash}"
  puts "\n\n" + green(">>>>>>>>>>>>>>>  BLUE/GREEN switch <<<<<<<<<<<<<<<")
  sleep 3

  puts blue("+++ STOP OLD containers")
  JSON.load(just_cmd("curl -s #{CONFIG["consul"]}/v1/catalog/services")).keys.select { |service|
    service != new_service && service =~ /^#{CONFIG["app"]}-\d+$/
  }.each do |service|
    JSON.load(just_cmd("curl -s #{CONFIG["consul"]}/v1/health/service/#{service}")).each do |s|
      do_it %(curl -s -X DELETE #{s["Node"]["Address"]}:8500/v1/agent/service/deregister/#{s["Service"]["ID"]})
    end
  end

  NODES.each do |n|
    keep_ids = just_cmd("ssh #{CONFIG["user"]}@#{n.ip} docker ps -q -f label=app=#{CONFIG["app"]} -f label=service=#{new_service}").split("\n")
    all_ids = just_cmd("ssh #{CONFIG["user"]}@#{n.ip} docker ps -q -f label=app=#{CONFIG["app"]}").split("\n")
    if (all_ids - keep_ids).length > 0
      do_it("ssh #{CONFIG["user"]}@#{n.ip} docker stop #{(all_ids - keep_ids).join(" ")}")
      do_it("ssh #{CONFIG["user"]}@#{n.ip} docker rm #{(all_ids - keep_ids).join(" ")}")
    end
  end

  puts blue("+++ CLEAN REGISTRY")
  (tags - [hash, current_hash]).each do |hash|
    digest = just_cmd("curl -s --head -H 'Accept: application/vnd.docker.distribution.manifest.v2+json' https://#{CONFIG["registry"]}/v2/#{CONFIG["app"]}/manifests/#{hash} | grep Docker-Content-Digest | cut -d' ' -f2")
    puts "digest = #{digest}"
    system "curl -X DELETE https://#{CONFIG["registry"]}/v2/#{CONFIG["app"]}/manifests/#{digest}"
  end

  puts green("Done.")
  if current_hash != "" && !target
    puts "to rollback execute: #{$0} #{ARGV[0]} deploy #{current_hash}"
  end
end


def print_status
  max_name_len = NODES.map { |n| n.name.length }.max
  max_ip_len = NODES.map { |n| n.ip.length }.max
  max_roles_len = NODES.map { |n| n.roles.inspect.length }.max
  current_service = just_cmd("curl -s #{CONFIG["consul"]}/v1/kv/apps/#{CONFIG["app"]}/service?raw")
  current_hash = just_cmd("curl -s #{CONFIG["consul"]}/v1/kv/apps/#{CONFIG["app"]}/hash?raw")
  current_dockers = NODES.map { |n| 
    just_cmd("ssh #{CONFIG["user"]}@#{n.ip} 'docker ps -q -f label=app=#{CONFIG["app"]} -f label=service=#{current_service} | wc -l'").to_i
  }
  dockers = NODES.map { |n| 
    just_cmd("ssh #{CONFIG["user"]}@#{n.ip} 'docker ps -q -f label=app=#{CONFIG["app"]} | wc -l'").to_i
  }
  puts ""
  puts green("App:     ") + CONFIG["app"]
  puts green("Service: ") + current_service
  puts green("Hash:    ") + current_hash
  NODES.each_with_index do |n, i|
    puts(
      " "*5 + n.name.rjust(max_name_len) + 
      " "*2 + n.ip.ljust(max_ip_len) + 
      " "*2 + n.roles.inspect.ljust(max_roles_len) +
      " "*2 + green(current_dockers[i]) + " / " + (dockers[i] - current_dockers[i] == 0 ? "0" : red(dockers[i] - current_dockers[i]))
    )
  end
end


def do_stop
  current_service = just_cmd("curl -s #{CONFIG["consul"]}/v1/kv/apps/#{CONFIG["app"]}/service?raw")
  if current_service != ""
    JSON.load(just_cmd("curl -s #{CONFIG["consul"]}/v1/health/service/#{current_service}")).each do |s|
      do_it %(curl -s -X DELETE #{s["Node"]["Address"]}:8500/v1/agent/service/deregister/#{s["Service"]["ID"]})
    end
  end
  NODES.map { |n| 
    ids = just_cmd("ssh #{CONFIG["user"]}@#{n.ip} docker ps -q -f label=app=#{CONFIG["app"]}").gsub("\n", " ")
    if ids != ""
      do_it "ssh #{CONFIG["user"]}@#{n.ip} docker stop #{ids}"
      do_it "ssh #{CONFIG["user"]}@#{n.ip} docker rm #{ids}"
    end
  }
  do_it "curl -X DELETE #{CONFIG["consul"]}/v1/kv/apps/#{CONFIG["app"]}?recurse"
  puts ""
end


if COMMAND == "deploy" || COMMAND == "d"
  deploy(ARGV[2])
elsif COMMAND == "status" || COMMAND == "s"
  print_status
elsif COMMAND == "stop"
  do_stop
else
  puts "FATAL: unknown command '#{COMMAND}'"
end