#!/usr/bin/env ruby
require 'mogilefs'
require 'optparse'
$stdin.binmode
$stdout.binmode
$stderr.sync = $stdout.sync = true

trap('INT') { exit 130 }
trap('PIPE') { exit 0 }
if md5_trailer_nodes = ENV["MD5_TRAILER_NODES"]
  md5_trailer_nodes.split(/\s*,\s*/).each do |host|
    MogileFS::HTTPFile::MD5_TRAILER_NODES[host] = true
  end
end

# this is to be compatible with config files used by the Perl tools
def parse_config_file!(path, dest = {})
  File.open(path).each_line do |line|
    case line
    when /^(domain|class)\s*=\s*(\S+)/
      dest[$1.to_sym] = $2
    when /^(?:trackers|hosts)\s*=\s*(.*)/
      dest[:hosts] = $1.split(/\s*,\s*/)
    when /^timeout\s*=\s*(.*)/
      dest[:timeout] = $1.to_f
    when /^noclobber\s*=\s*true\s*/
      dest[:noclobber] = true
    else
      warn "Ignored configuration line: #{line}" unless /^#/.match(line)
    end
  end
  dest
end

# parse the default config file if one exists
def_file = File.expand_path("~/.mogilefs-client.conf")
def_cfg = File.exist?(def_file) ? parse_config_file!(def_file) : {}

# parse the command-line first, these options take precedence over all else
cli_cfg = {}
config_file = nil
ls_l = false
ls_h = false
chunk = false
range = false
test = {}
cat = { :raw => false }

ARGV.options do |x|
  x.banner = "Usage: #{$0} [options] <command> [<arguments>]"
  x.separator ''

  x.on('-c', '--config=/path/to/config',
       'config file to load') { |file| config_file = file }

  x.on('-t', '--trackers=host1[,host2]', '--hosts=host1[,host2]', Array,
       'hostnames/IP addresses of trackers') do |trackers|
    cli_cfg[:hosts] = trackers
  end

  x.on('-e', 'True if key exists') { test[:e] = true }
  x.on('-r', '--raw', 'show raw big_info file information') { cat[:raw] = true }
  x.on('-n', '--no-clobber', 'do not clobber existing key') do
    cli_cfg[:noclobber] = true
  end

  x.on('-C', '--class=s', 'class') { |klass| cli_cfg[:class] = klass }
  x.on('-d', '--domain=s', 'domain') { |domain| cli_cfg[:domain] = domain }
  x.on('-l', "long listing format (`ls' command)") { ls_l = true }
  x.on('-h', '--human-readable',
       "print sizes in human-readable format (`ls' command)") { ls_h = true }
  x.on('--chunk', "chunk uploads (`tee' command)") { chunk = true }
  x.on('--range', "stream partial uploads (`tee' command)") { range = true }
  x.separator ''
  x.on('--help', 'Show this help message.') { puts x; exit }
  x.on('--version', 'Show --version') { puts "#$0 #{MogileFS::VERSION}"; exit }
  x.parse!
end

# parse the config file specified at the command-line
file_cfg = config_file ? parse_config_file!(config_file) : {}

# read environment variables, too.  This Ruby API favors the term
# "hosts", however upstream MogileFS teminology favors "trackers" instead.
# Favor the term more consistent with what the MogileFS inventors used.
env_cfg = {}
if ENV["MOG_TRACKERS"]
  env_cfg[:hosts] = ENV["MOG_TRACKERS"].split(/\s*,\s*/)
end
if ENV["MOG_HOSTS"] && (env_cfg[:hosts] || []).empty?
  env_cfg[:hosts] = ENV["MOG_HOSTS"].split(/\s*,\s*/)
end
env_cfg[:domain] = ENV["MOG_DOMAIN"] if ENV["MOG_DOMAIN"]
env_cfg[:class] = ENV["MOG_CLASS"] if ENV["MOG_CLASS"]

# merge the configs, favoring them in order specified:
cfg = {}.merge(def_cfg).merge(env_cfg).merge(file_cfg).merge(cli_cfg)

# error-checking
err = []
err << "trackers must be specified" if cfg[:hosts].nil? || cfg[:hosts].empty?
err << "domain must be specified" unless cfg[:domain]
if err.any?
  warn "Errors:\n  #{err.join("\n  ")}"
  warn ARGV.options
  exit 1
end

unless cmd = ARGV.shift
  warn ARGV.options
  exit 1
end

cfg[:timeout] ||= 30 # longer timeout for interactive use
mg = MogileFS::MogileFS.new(cfg)

def store_file_retry(mg, key, storage_class, filepath)
  tries = 0
  begin
    mg.store_file(key, storage_class, filepath)
  rescue MogileFS::UnreadableSocketError,
         MogileFS::Backend::NoDevicesError => err
    if ((tries += 1) < 10)
      warn "Retrying on error: #{err}: #{err.message} tries: #{tries}"
      retry
    else
      warn "FATAL: #{err}: #{err.message} tries: #{tries}"
    end
    exit 1
  end
end

def human_size(size)
  suff = ''
  %w(K M G).each do |s|
    size /= 1024.0
    if size <= 1024
      suff = s
      break
    end
  end
  sprintf("%.1f%s", size, suff)
end

begin
  case cmd
  when 'cp'
    filename = ARGV.shift or raise ArgumentError, '<filename> <key>'
    dkey = ARGV.shift or raise ArgumentError, '<filename> <key>'
    ARGV.shift and raise ArgumentError, '<filename> <key>'
    cfg[:noclobber] && mg.exist?(dkey) and
      abort "`#{dkey}' already exists and -n/--no-clobber was specified"
    store_file_retry(mg, dkey, cfg[:class], filename)
  when 'cat'
    ARGV.empty? and raise ArgumentError, '<key1> [<key2> ...]'
    ARGV.each do |key|
      if (!cat[:raw] && key =~ /^_big_info:/)
        mg.bigfile_write(key, $stdout, {:verify => true})
      else
        mg.get_file_data(key, $stdout)
      end
    end
  when 'updateclass'
    newclass = cfg[:class] or abort '-C/--class not specified'
    ARGV.each do |key|
      mg.updateclass(key, newclass)
    end
  when 'ls'
    prefixes = ARGV.empty? ? [ nil ] : ARGV
    if ls_l
      each_key = lambda do |key, size, devcount|
        size = ls_h && size > 1024 ? human_size(size) : size.to_s
        size = (' ' * (12 - size.length)) << size # right justify
        puts [ sprintf("% 2d", devcount), size, key ].pack("A4 A16 A*")
      end
    else
      each_key = lambda { |key| puts key }
    end
    prefixes.each { |prefix| mg.each_key(prefix, &each_key) }
  when 'rm'
    ARGV.empty? and raise ArgumentError, '<key1> [<key2>]'
    ARGV.each { |key| mg.delete(key) }
  when 'mv'
    from = ARGV.shift or raise ArgumentError, '<from> <to>'
    to = ARGV.shift or raise ArgumentError, '<from> <to>'
    ARGV.shift and raise ArgumentError, '<from> <to>'
    mg.rename(from, to)
  when 'stat' # this outputs a RFC822-like format
    ARGV.empty? and raise ArgumentError, '<key1> [<key2>]'
    ok = true
    ARGV.each_with_index do |key,j|
      begin
        info = mg.file_info(key)
        puts "Key: #{key}"
        puts "Size: #{info['length']}"
        puts "Class: #{info['class']}"
        checksum = info['checksum'] and puts "Checksum: #{checksum}"
        o = { :pathcount => info["devcount"] }
        mg.get_paths(key, o).each_with_index do |path,i|
          puts "URL-#{i}: #{path}"
        end
        puts "" if ARGV.size != (j + 1)
      rescue MogileFS::Backend::UnknownKeyError
        warn "No such key: #{key}"
        ok = false
      end
    end
    exit(ok)
  when 'tee'
    abort "--range and --chunk are incompatible" if range && chunk
    dkey = ARGV.shift or raise ArgumentError, '<key>'
    ARGV.shift and raise ArgumentError, '<key>'
    cfg[:noclobber] && mg.exist?(dkey) and
      abort "`#{dkey}' already exists and -n/--no-clobber was specified"
    skip_tee = File.stat('/dev/null') == $stdout.stat
    largefile = :tempfile
    largefile = :content_range if range
    largefile = :stream if chunk

    io = mg.new_file(dkey, :class => cfg[:class], :largefile => largefile)
    begin
      buf = $stdin.readpartial(16384)
      begin
        io.write(buf)
        $stdout.write(buf) unless skip_tee
        $stdin.readpartial(16384, buf)
      end while true
    rescue EOFError
    end
    io.close
  when 'test'
    truth, ok = true, nil
    raise ArgumentError, "-e must be specified" unless (test.size == 1)

    truth, key = case ARGV.size
    when 1
      [ true, ARGV[0] ]
    when 2
      if ARGV[0] != "!"
        raise ArgumentError, "#{ARGV[0]}: binary operator expected"
      end
      [ false, ARGV[1] ]
    else
      raise ArgumentError, "Too many arguments"
    end

    test[:e] or raise ArgumentError, "Unknown flag: -#{test.keys.first}"
    ok = mg.exist?(key)
    truth or ok = ! ok
    exit ok ? 0 : 1
  else
    raise ArgumentError, "Unknown command: #{cmd}"
  end
rescue ArgumentError => err
  warn "Usage: #{$0} #{cmd} #{err.message}"
  exit 1
end
exit 0