#!/usr/bin/env ruby

require 'rubygems'
require 'log4r'
require 'yaml'
require 'shellwords'
require 'optparse'
require 'net/ssh'

# Based on actions available to the VirtualBox provider:
#   https://github.com/mitchellh/vagrant/tree/master/plugins/providers/virtualbox
class Provider
  WAIT = 5

  def initialize(config)
    @config = config
    @logger = Log4r::Logger.new("vagrant::provider::lxc")
    @logger.outputters = Log4r::Outputter.stdout
    if config['output']
      @logger.outputters << Log4r::FileOutputter.new('output', 'filename' => config['output'])
    end
    # @logger.level = Log4r::INFO
  end

  # @see Vagrant::Plugin::V1::Provider#action
  def action(name, *args)
    # Attempt to get the action method from this class if it
    # exists, otherwise return nil to show that we don't support the
    # given action.
    action_method = "action_#{name}"
    return send(action_method, *args) if respond_to?(action_method)
    nil
  end

  protected

  def run(cmd)
    @logger.debug "Running: #{cmd}"
    system cmd
  end

  def action_up
    was_created = container_created?
    if was_created
      @logger.info("Container already created, moving on...")
    else
      @logger.info("Creating container...")
      # TRY: run 'sudo lxc-create -t ubuntu -n vagrant-container -b vagrant'
      # DISCUSS: Copy key directly to /var/lib/lxc/$host/root/.ssh/authorized_keys to be generic?
      unless run 'sudo lxc-create -t ubuntu-cloud -n vagrant-container -- -S /home/vagrant/.ssh/id_rsa.pub'
        puts 'Error creating box'
        exit 1
      end
      unless container_created?
        puts 'Error creating container'
        exit 1
      end
    end

    if container_started?
      @logger.info('Container already started')
    else
      share_folders

      @logger.info('Starting container...')
      unless run "sudo lxc-start -n vagrant-container -d #{configs}"# -o /tmp/lxc-start.log -l DEBUG"
        puts 'Error starting container!'
        exit 1
      end
      run 'sudo lxc-wait --name vagrant-container --state RUNNING'
      unless container_started?
        puts 'Error starting container!'
        exit 1
      end
      @logger.info('Container started')

      forward_ports

      unless was_created
        @logger.debug "Waiting #{WAIT} seconds before setting up vagrant user"
        sleep WAIT
        setup_vagrant_user
      end
    end
  end

  def action_halt
    if container_started?
      @logger.info('Stopping container...')
      unless run 'sudo lxc-shutdown -n vagrant-container'
        puts 'Error halting container!'
        exit 1
      end
      run 'sudo lxc-wait --name vagrant-container --state STOPPED'
      if container_started?
        puts 'Error halting container!'
        exit 1
      end
      @logger.info('Container halted')
    else
      @logger.info('Container already halted')
    end
  end

  def action_destroy
    if container_created?
      if container_started?
        action_halt
        @logger.debug "Waiting #{WAIT} seconds to proceed with destroy..."
        sleep WAIT
      end
      @logger.info("Destroying container...")
      unless run 'sudo lxc-destroy -n vagrant-container'
        puts 'Error destroying container'
        exit 1
      end
      if container_created?
        puts 'Error destroying container'
        exit 1
      end
      @logger.debug "Waiting #{WAIT} seconds for things to settle down..."
      sleep WAIT
      @logger.info("Container destroyed")
    else
      @logger.info("Container not created")
    end
  end

  def action_reload
    action_halt if container_started?
    action_up
  end

  # TODO: Switch over to Net:SSH
  def action_ssh(opts = {'user' => 'vagrant'})
    # FIXME: We should not depend on an IP to be configured
    raise 'SSH support is currently available to a predefined IP only' unless @config['ip']

    cmd = "ssh #{opts['user']}@#{@config['ip']} -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o LogLevel=quiet"
    cmd << " -- #{Shellwords.escape opts['command']}" if opts['command']

    unless run(cmd)
      puts 'Error running ssh command!'
      exit 1
    end
  end

  def setup_vagrant_user
    unless @config['ip']
      # FIXME: Need to find a way to grab the container IP
      @logger.warn('Unfortunately automatic vagrant user setup does not work unless an IP is specified')
      return
    end

    @logger.info 'Setting up vagrant user'

    # TODO: We could try to use lxc-attach instead of SSH

    # Based on:
    #   https://github.com/jedi4ever/veewee/blob/master/templates/ubuntu-12.10-server-amd64-packages/vagrant.sh
    cmds = [
      #'groupadd -r admin',
      'useradd -d /home/vagrant -m vagrant -r -s /bin/bash',
      'usermod -a -G admin vagrant',
      'cp /etc/sudoers /etc/sudoers.orig',
      'sed -i -e \'/Defaults\s\+env_reset/a Defaults\texempt_group=admin\' /etc/sudoers',
      'sed -i -e \'s/%admin\s\+ALL=(ALL)\s\+ALL/%admin ALL=NOPASSWD:ALL/g\' /etc/sudoers',
      'service sudo restart',
      '-u vagrant -- mkdir -p /home/vagrant/.ssh',
      '-u vagrant -- curl -o /home/vagrant/.ssh/authorized_keys https://raw.github.com/mitchellh/vagrant/master/keys/vagrant.pub'
    ]

    # FIXME: Needs to abort the process if any of this commands fail
    ssh_conn('ubuntu') do |ssh|
      cmds.each do |cmd|
        @logger.debug "SSH: sudo #{cmd}"
        ssh.exec!("sudo #{cmd}")
      end
    end
  end

  def ssh_conn(user = 'vagrant')
    Net::SSH.start(@config['ip'], user, :user_known_hosts_file => '/dev/null') do |ssh|
      yield ssh
    end
  end

  def container_created?
    `lxc-ls` =~ /^vagrant\-container/
  end

  def container_started?
    `sudo -- lxc-info -n vagrant-container` =~ /RUNNING/
  end

  def share_folders
    @logger.info('Setting up shared folders...')

    mount_folder(File.expand_path('.'), '/vagrant')

    Array(@config['shared_folders']).each do |folder|
      mount_folder(folder['source'], folder['destination'])
    end
  end

  def mount_folder(source, destination)
    @logger.info("Sharing #{source} as #{destination}")
    run <<STR
sudo mount --bind #{source} #{source}
sudo mount --make-unbindable #{source}
sudo mount --make-shared #{source}

if ! [ -d /var/lib/lxc/vagrant-container/rootfs#{destination} ]; then
  sudo mkdir -p /var/lib/lxc/vagrant-container/rootfs#{destination}
fi

if ! $(sudo grep -q '#{source} /var/lib/lxc/vagrant-container/rootfs#{destination}' /var/lib/lxc/vagrant-container/fstab); then
  cat <<EOF | sudo tee -a /var/lib/lxc/vagrant-container/fstab
#{source} /var/lib/lxc/vagrant-container/rootfs#{destination} none bind 0 0
EOF
fi
STR
  end

  def configs
    configs = []
    configs << "-s lxc.network.ipv4='#{@config['ip']}'" if @config['ip']
    configs << '-s lxc.cgroup.memory.limit_in_bytes=400M'
    configs << '-s lxc.cgroup.memory.memsw.limit_in_bytes=500M'
    configs.join(' ')
  end

  def forward_ports
    return unless @config.key?('forwards')

    @logger.info('Forwarding ports...')
    forwards = ''
    @config['forwards'].each do |forward|
      host_port, guest_port = forward
      @logger.info("-- #{guest_port} => #{host_port}")
      forwards << "0.0.0.0 #{host_port} #{@config['ip']} #{guest_port}"
    end

    # FIXME: We should be nice to others and not overwrite the config all the time ;)
    File.open('/etc/rinetd.conf', 'w') do |f|
      f.puts forwards
      f.puts 'logfile /var/log/rinetd.log'
    end
    @logger.info('Restarting rinetd')
    `sudo service rinetd restart`
  end
end

raise 'You need to provide an action' unless ARGV[0]

action = ARGV.shift.to_sym
if action == :ssh
  options = {'user' => 'vagrant'}
  OptionParser.new do |opts|
    opts.on("-c", '--command [COMMAND]') { |v| options['command'] = v }
    opts.on('-u', '--user [USER]') { |v| options['user'] = v }
  end.parse!
  arguments = [options]
else
  init_options = {}
  OptionParser.new do |opts|
    opts.on("-o", '--output [FILE]') { |v| init_options['output'] = v }
  end.parse!
end

config = YAML.load File.open('./config.yml') if File.exists? './config.yml'
config ||= {}
config['output'] = init_options.delete('output') if init_options && init_options.key?('output')

@provider = Provider.new(config || {})

@provider.action(action, *(arguments || []))