#!/usr/bin/env ruby

require "open3"
require 'logger'
require 'yaml'
require 'optparse'
require 'fileutils'


class MultiIO
  def initialize(*targets)
     @targets = targets
  end

  def write(*args)
    @targets.each {|t| t.write(*args)}
  end

  def close
    @targets.each(&:close)
  end
end


def windows?
  !!(RbConfig::CONFIG['host_os'] =~ /mingw/)
end

def logpath
  log_folder = windows? ? "C:/RCS/DB/log" : File.expand_path("../../log", __FILE__)

  if Dir.exists?(log_folder)
    "#{log_folder}/#{File.basename(__FILE__)}.log"
  else
    __FILE__+".log"
  end
end

def logger
  @logger ||= begin
    if $options[:only_stdout]
      lgr = Logger.new(STDOUT)
    else
      log_file = File.open(logpath, "a")
      puts "Output will also goes to #{logpath}"
      lgr = Logger.new MultiIO.new(STDOUT, log_file)
    end

    lgr.level = $options[:debug] ? Logger::DEBUG : Logger::INFO
    lgr.formatter = proc { |severity, datetime, progname, msg| "#{datetime} [#{severity.rjust(5)}]: #{msg}\n" }
    lgr
  end
end

def mongo_eval(address, command, bin: nil)
  bin ||= windows? ? "C:/RCS/DB/mongodb/win/mongo.exe" : "mongo"
  bin << ".exe" if windows? and !bin.end_with?(".exe")
  run("#{bin} #{address} --eval \"#{command}\"")
end

def cn
  path = windows? ? "C:/RCS/DB/config/config.yaml" : File.expand_path("../../config/config.yaml", __FILE__)
  YAML.load_file(path)["CN"]
end

def mongo_version
  mongo_eval("test", "db.runCommand({buildinfo: 1}).version").last.strip
end

def old_mongo?
  mongo_version.start_with?($options[:from])
end

def new_mongo?
  mongo_version.start_with?($options[:to])
end

def mongo_start_balancer
  mongo_eval("admin", "sh.startBalancer()")
end

def mongo_balancer_stopped?
  mongo_eval("config", "db.settings.find({_id: 'balancer'})[0].stopped").last.strip == 'true'
end

def mongo_stop_balancer
  mongo_eval("admin", "sh.stopBalancer()")
end

def options
  $options
end

# @note Order of captured stderr may be different.
# Better solution is to use pty but it does not work under Windows.
def run(command, return_exit_status: false)
  output = []
  logger.debug(command)
  Open3.popen2e(command) do |stdin, stdout_err, wait_thr|
    while line = stdout_err.gets
      line.strip!
      logger.debug("[stdouterr] #{line}")
      output << line.strip
    end

    return wait_thr.value.success? if return_exit_status
  end
  logger.debug("popen finished")
  output
end

def check_version!(binary, version)
  binary << ".exe" if windows? and !binary.end_with?(".exe")
  raise "File not found #{binary}" unless File.exists?(binary)
  raise "#{binary}: expected version #{version}" unless run("#{binary} --version").first =~ /version #{version}/
end

def mongo_upgrade
  bin = "#{$options[:new_mongo_dir]}/mongos"
  bin << ".exe" if windows?

  check_version!(bin, $options[:to])

  output = run("#{bin} --configdb \"#{cn}\" --upgrade")

  if !$options[:redo] and !$options[:shard] and !output.find { |line| line =~ /upgrade of config server to v5 successful/ }
    raise "upgrade terminated unsuccessfully"
  end

  if !output.find { |line| line =~ /Config database is at version v5/ }
    raise "upgrade failed"
  end
end

def mongo_bin_path
  windows? ? "C:/RCS/DB/mongodb/win" : File.expand_path("../../mongodb/macos", __FILE__)
end

def mongo_shutdown
  %w[27017 27019 27018].each do |port|
    next if $options[:shard] and port == '27019'
    address = "127.0.0.1:#{port}"
    output = mongo_eval(address+"/admin", "db.shutdownServer()")
    logger.warn("Unable to shutdown mongoDB at #{address}") unless output.last =~ /(server should be down|connect failed)/
  end

  if windows?
    windows_service(:stop, "RCS Master Router", force: true)
    windows_service(:stop, "RCS Master Config", force: true)
    windows_service(:stop, "RCS Shard", force: true)

    waiting = 0

    loop do
      list = run("tasklist").join(" ") rescue ""
      break if list !~ /mongo(d|s).exe/i
      sleep(1)
      waiting += 1

      if waiting >= 10
        logger.warn("mongod or mongos are still running, killing...")
        run("taskkill /F /IM mongos.exe")
        run("taskkill /F /IM mongod.exe")
        break
      end
    end
  end
end

def mongo_startup
  if windows?
    ["RCS Shard", "RCS Master Config", "RCS Master Router"].each do |name|
      next if $options[:shard] and name == "RCS Master Config"
      windows_service(:start, name)
    end
  else
    script = File.expand_path("../../bin/rcs-db-mongo-all", __FILE__)
    run(script)
  end
end

def is_mongo_up?(wait_until_up: 0)
  waited_time = 0

  loop do
    is_up = mongo_eval("admin", "1").last == "1"
    return true if is_up
    break if waited_time >= wait_until_up
    sleep(1)
    waited_time += 1
  end

  return false
end

def replace_binaries
  Dir["#{mongo_bin_path}/*"].each { |path| FileUtils.rm(path) }
  Dir["#{$options[:new_mongo_dir]}/*"].each { |path| FileUtils.cp(path, mongo_bin_path) }
end

def windows_service(action, name, force: false)
  cmd = "NET #{action.to_s.upcase} \"#{name}\""
  cmd << " /y" if force

  success = run(cmd, return_exit_status: true)

  unless success
    logger.warn("Failed to #{action} service #{name}")
  end
end

def basic_upgrade_check
  dbs = mongo_eval("config", "db.databases.distinct('_id')").last

  logger.debug("Databases are #{dbs}")

  mongo_eval("config", "db.changelog.ensureIndex({_id: 1}, {unique: true})")

  upgrade_check('config')
end

def upgrade_check(db_name)
  if mongo_eval(db_name, "db.upgradeCheck()", bin: "#{$options[:new_mongo_dir]}/mongo").last != "true"
    raise "upgradeCheck() failed on database #{db_name}"
  end
end

def same_folder?(a, b)
  File.expand_path(a).downcase.gsub("\\", "/") == File.expand_path(b).downcase.gsub("\\", "/")
end

ARGV << '--help' if ARGV.empty?

$options = {}
$options[:from] = '2.4'
$options[:to]   = '2.6'

optparse = OptionParser.new do |parser|
  parser.banner << "\nUpgrade mongoDB from version #{$options[:from]}.x to version #{$options[:to]}"
  parser.banner << "\nIt is advisable that all the RCS services (except for the mongoDB related) are stopped while running this."

  parser.on('-d', '--debug', 'Set logger level to DEBUG (instead of INFO)') {
    $options[:debug] = true
  }

  parser.on("-p", "--mongo#{$options[:to].gsub('.', '')}-dir PATH", String, "Path of the folder with all the mongoDB #{$options[:to]} binaries") { |path|
    $options[:new_mongo_dir] = File.expand_path(path.gsub('\\', '/'))
  }

  parser.on('--safe-redo', 'Run even if the current version is already (partially) upgraded') {
    $options[:redo] = true
  }

  parser.on('--only-stdout', 'Output goes only to STDOUT. Default is to STDOUT and to a logfile') {
    $options[:only_stdout] = true
  }

  parser.on('--upgrade-check DATABASE', 'Run the upgradeCheck function on DATABASE') { |db_name|
    $options[:upgrade_check] = db_name
  }

  parser.on('--shard', 'Use this option when upgrading a shard') {
    $options[:shard] = true
  }
end.parse!

begin
  $stdout.sync = true

  logger.info "Check if mongoDB is up"

  unless is_mongo_up?
    logger.info "mongoDB is down. Starting mongoDB services"
    mongo_startup
  end

  logger.info "Check if current mongoDB bin folder is valid"
  raise "Missing folder #{mongo_bin_path}" unless Dir.exists?(mongo_bin_path)

  logger.info "Check if given mongoDB #{$options[:to]} bin folder is valid"
  check_version!("#{$options[:new_mongo_dir]}/mongos", $options[:to])
  raise "mongoDB #{$options[:to]} bin folder must not be #{mongo_bin_path}" if same_folder?($options[:new_mongo_dir], mongo_bin_path)

  logger.info "Current mongoDB version is #{mongo_version}"
  logger.info "Current cn is #{cn}"

  if !$options[:redo]
    if new_mongo?
      logger.info "mongoDB #{$options[:to]} is already installed"
      exit(0)
    end

    raise "This upgrade require mongoDB version #{$options[:from]}" unless old_mongo?
  end

  if $options[:upgrade_check]
    upgrade_check($options[:upgrade_check])
    exit!
  end

  logger.info "Ensuring upgrade prerequisites"

  basic_upgrade_check

  if mongo_balancer_stopped?
    logger.warn "Balancer is already stopped"
  else
    logger.info "Stopping balancer"

    mongo_stop_balancer

    if mongo_balancer_stopped?
      logger.info "balancer stopped"
    else
      raise "Unable to stop balancer"
    end
  end

  logger.info "Starting upgrade"

  mongo_upgrade

  logger.info "Upgrade completed"

  logger.info "Shutting down mongoDB"

  mongo_shutdown

  raise "Shutdown failed. mongoDB is still up" if is_mongo_up?

  logger.info "Replacing binaries"

  replace_binaries

  logger.info "Staring up mongoDB"

  mongo_startup

  logger.info "Checking if mongoDB is up"

  raise "Statup failed" unless is_mongo_up?(wait_until_up: 60)

  logger.info "Restarting the balancer"

  mongo_start_balancer

  raise "Unable to start the balancer" if mongo_balancer_stopped?

  if File.exists?("#{$options[:new_mongo_dir]}/mongo.exe")
    logger.info "Removing new binaries dir"
    logger.debug "rm -rf #{$options[:new_mongo_dir]}"
    FileUtils.rm_rf("#{$options[:new_mongo_dir]}")
  end

  logger.info "Upgrade completed successfully :)"
rescue SystemExit => e
  exit(e.status)
rescue Interrupt
  logger.error("Interrupted")
  exit(1)
rescue Exception => e
  logger.error("[#{e.class}] #{e.message.inspect} @ #{e.backtrace[0].strip if e.backtrace}")
  raise(e)
  exit(1)
end