#!/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 = {}, [], [], {}, {}

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

    cluster = @clusters[c]

    if cluster
        cluster_hosts = cluster["hosts"]

        @i2_options[:login_override] = !cluster["login"].nil? ? cluster["login"] : @i2_options[:login_override]
        @i2_options[:login_override] = login_from_cli if login_from_cli
        @i2_options[:broadcast] = !cluster["broadcast"].nil? ? cluster["broadcast"] : @i2_options[:broadcast]
        @i2_options[:profile] = !cluster["profile"].nil? ? cluster["profile"]  : @i2_options[:profile]
        @i2_options[:rank] = !cluster["rank"].nil? ? cluster["rank"] : @i2_options[:rank]

        if @i2_options[:login_override] then
            cluster_hosts = cluster_hosts.map{|h| "#{@i2_options[:login_override]}@#{h}"}
        end

        @ssh_environment.merge!(cluster["environment"].inject({}){|m, v| m.merge(v)}) if cluster["environment"]

        @servers += cluster_hosts
    else
        puts "ERROR: unknown cluster #{c}. Check your #{@config_file}"
        exit 1
    end
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
        @clusters = config_hash["clusters"]
        
        # Options from the config file
        @i2_options[:iterm2] = config_hash["iterm2"]
        @i2_options[:login_override] = config_hash["login"]
        @i2_options[:broadcast] = config_hash["broadcast"]
        @i2_options[:profile] = config_hash["profile"]
        @i2_options[:rank] = config_hash["rank"]
        
        @ssh_environment.merge!(config_hash["environment"].inject({}){|m, v| m.merge(v)}) if config_hash["environment"]
        
    else
        # Convert version 1 format to version 2
        clusters = config_hash["clusters"].inject({}){|m, c| m[c[0]] = {"hosts" => c[1]}; m}
    end

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|
        @i2_options[:login_override] = 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
        @i2_options[:rank] = true
    end

    # iTerm2 options
    opts.on '-F', '--fullscreen',
        'Make the window fullscreen' do
        @i2_options[:fullscreen] = true
    end
    opts.on '-C', '--columns COLUMNS', Integer,
        'Number of columns (rows will be calculated)' do |c|
        if @i2_options[:rows]
            puts "ERROR: -C and -R can't be used at the same time"
            puts optparse.help
            exit 1
        else
            @i2_options[:columns] = c
        end
    end
    opts.on '-R', '--rows ROWS', Integer,
        'Number of rows (columns will be calculated)' do |r|
        if @i2_options[:columns]
            puts "ERROR: -C and -R can't be used at the same time"
            puts optparse.help
            exit 1
        else
            @i2_options[:rows] = r
        end
    end
    opts.on '-b', '--broadcast',
        'Start with broadcast input (DANGEROUS!)' do
        @i2_options[:broadcast] = true
    end
    opts.on '-nb', '--nobroadcast',
        'Disable broadcast' do
        @i2_options[:broadcast] = false
    end
    opts.on '-p', '--profile PROFILE',
        'Name of the iTerm2 profile (default: Default)' do |p|
        @i2_options[:profile] = p
    end
    opts.on "-2", '--iterm2',
        'Use iTerm2 instead of iTerm' do
        @i2_options[:iterm2] = true
    end

end
optparse.parse!

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

if @i2_options[:login_override] then
    @servers = @servers.map{|h| "#{@i2_options[:login_override]}@#{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