#!/usr/bin/ruby
require 'rubygems'
require 'optparse'
require 'i2cssh'
require 'yaml'

@config_file = File.expand_path "~/.i2csshrc"

@i2_options, ssh_options, @servers, @clusters, @ssh_environment, opts_from_cmdline = {}, [], [], {}, {}, {}

def get_hosts(c)
    if c =~ /(.+)@(.+)/ then
        login_from_cli = $1
        c = $2
    end

    cluster = @clusters[c]

    if cluster
        set_options(cluster, login_from_cli)

        if @i2_options[:login] then
            @servers += cluster["hosts"].map{|h| "#{@i2_options[:login]}@#{h}"}
        else
            @servers += cluster["hosts"]
        end
    else
        puts "ERROR: unknown cluster #{c}. Check your #{@config_file}"
        exit 1
    end
end

def set_options(config_hash, login_override=nil)
    if config_hash["columns"] and config_hash["rows"] then
        puts "CONFIG ERROR: rows and columns can't be used at the same time"
        exit 1
    end

    [:broadcast, :profile, :rank, :iterm2, :login, :columns, :rows, :sleep, :direction, :itermname].each do |p|
        @i2_options[p] = config_hash[p.to_s].nil? ? @i2_options[p] : config_hash[p.to_s]
    end

    @i2_options[:login] = login_override if login_override
    @i2_options[:direction] ||= :column
    @i2_options[:direction] = @i2_options[:direction].to_sym
    @ssh_environment.merge!(config_hash["environment"].inject({}){|m, v| m.merge(v)}) if config_hash["environment"]
end

if File.exists?(@config_file) 
    config_hash = YAML.load File.read @config_file

    # Read config and set defaults from config
    if config_hash["version"] && config_hash["version"].to_i >= 2 then
        set_options(config_hash)
        @clusters = config_hash["clusters"]
    else
        # Convert version 1 format to version 2
        @clusters = config_hash["clusters"].inject({}){|m, c| m[c[0]] = {"hosts" => c[1]}; m}
    end

else
    set_options({})
end

optparse = OptionParser.new do |opts|
    opts.banner = "Usage: #{File.basename(__FILE__)} [options] [(username@host [username@host] | username@cluster)]"

    # Check if we have a cluster. 
    opts.on '-c', '--cluster CLUSTERNAME',
        'Name of the cluster specified in ~/.i2csshrc' do |c|
        get_hosts(c)
    end

    opts.on '-m', '--machines a,b,c', Array,
        'Comma-separated list of hosts' do |h|
        @servers += h
    end

    # Hosts
    opts.on '-f', '--file FILE',
        'Cluster file (one hostname per line)' do |f|
        @servers += File.read(f).split "\n"
    end

    # Command line options override config file

    # SSH options
    opts.on '-A', '--forward-agent',
        'Enable SSH agent forwarding' do
        ssh_options << '-A'
    end
    opts.on '-l', '--login LOGIN',
        'SSH login name' do |u|
        opts_from_cmdline[:login] = u
        
    end
    opts.on '-e', '--environment KEY=VAL',
        'Send environment vars (comma-separated list, need to start with LC_)' do |e|
        @ssh_environment = e.split(",").inject({}) {|m, x| key, val = x.split("="); m[key] = val; m}
    end
    opts.on '-r', '--rank',
        'Send LC_RANK with the host number as environment variable' do
        opts_from_cmdline[:rank] = true
    end

    # iTerm2 options
    opts.on '-F', '--fullscreen',
        'Make the window fullscreen' do
        opts_from_cmdline[:fullscreen] = true
    end
    opts.on '-C', '--columns COLUMNS', Integer,
        'Number of columns (rows will be calculated)' do |c|
        if opts_from_cmdline[:rows]
            puts "ERROR: -C and -R can't be used at the same time"
            puts optparse.help
            exit 1
        else
            opts_from_cmdline[:columns] = c
        end
    end
    opts.on '-R', '--rows ROWS', Integer,
        'Number of rows (columns will be calculated)' do |r|
        if opts_from_cmdline[:columns]
            puts "ERROR: -C and -R can't be used at the same time"
            puts optparse.help
            exit 1
        else
            opts_from_cmdline[:rows] = r
        end
    end
    opts.on '-b', '--broadcast',
        'Start with broadcast input (DANGEROUS!)' do
        opts_from_cmdline[:broadcast] = true
    end
    opts.on '-nb', '--nobroadcast',
        'Disable broadcast' do
        opts_from_cmdline[:broadcast] = false
    end
    opts.on '-p', '--profile PROFILE',
        'Name of the iTerm2 profile (default: Default)' do |p|
        opts_from_cmdline[:profile] = p
    end
    opts.on "-2", '--iterm2',
        'Use iTerm2 instead of iTerm' do
        opts_from_cmdline[:iterm2] = true
    end
    opts.on "-i", '--itermname NAME', String,
        'Name of the application to use (default: iTerm)' do |i|
        opts_from_cmdline[:itermname] = i
    end
    opts.on '-s', '--sleep SLEEP', Float,
        'Number of seconds to sleep between creating SSH sessions' do |s|
        opts_from_cmdline[:sleep] = s
    end
    opts.on "-d", '--direction DIRECTION', String,
        'Direction that new sessions are created (default: column)' do |d|
        unless ["row", "column"].include?(d) then
            puts "ERROR: -d requires 'row' or 'column'"
            puts optparse.help
            exit 1
        end
        opts_from_cmdline[:direction] = d.to_sym
    end
    opts.on '-X', '--extra EXTRA_PARAM', String,
        'Additional ssh parameters (e.g. -Xi=myidentity.pem)' do |x|

         ssh_options << "-" + x.split("=").join(" ")
    end

end
optparse.parse!

if ARGV.length == 1 then
    c = ARGV[0]
    get_hosts(c)
elsif ARGV.length > 1 then
    @servers = ARGV
end

@i2_options.merge!(opts_from_cmdline)

if @i2_options[:login] then
    @servers = @servers.map{|h| "#{@i2_options[:login]}@#{h.gsub(/.+@/,'')}"}
end

if @servers.empty?
    puts "ERROR: no servers given"
    puts optparse.help
    exit
end

I2Cssh.new @servers, ssh_options, @i2_options, @ssh_environment