#!/usr/bin/env ruby

require 'optparse'
require 'tmpdir'

@options = {}
OptionParser.new do |opts|
  opts.banner = <<-BANNER.gsub(/^    /,'')
    Box : manage box

    Usage: #{File.basename($0)} boot [options] 

    Options are:
    BANNER
  opts.separator ""
  opts.on("-r", "--root=disk|iso", String,
          "The support used to boot") { |arg| @options[:root] = arg }
  opts.on("-t", "--target=<disk device>", String,
          "The target to clone an image") { |arg| @options[:target] = arg }
  opts.on("-i", "--image=<disk image>", String,
          "The box image to be cloned") { |arg| @options[:image] = arg }
  opts.on("-o", "--host=<host>", String,
          "The host to be upgraded") { |arg| @options[:host] = arg }
  opts.on("-s", "--skip-format",
          "Skip partitionning and formatting") { |arg| @options[:skip_format] = true }
  opts.on("-e", "--external-extlinux",
          "Don't use extlinux provided by image") { |arg| @options[:external_extlinux] = true }
  opts.on("-b", "--name=<box name>", String,
          "The box name") { |arg| @options[:name] = arg }
  opts.on("-m", "--mac-index=<index>", Integer,
          "An numeric index uses to choose mac") { |arg| @options[:mac_index] = arg }
  opts.on("-n", "--dry-run",
          "Dry run mode") { |arg| @options[:dry_run] = true }
  opts.on("-h", "--help",
          "Show this help message.") { puts opts; exit }
  opts.parse!(ARGV)
  
  @command = ARGV.shift

  unless %w{boot clone upgrade}.include? @command
    puts opts; exit
  end
end

module SystemTools

  def sh(command)
    puts "* #{command}"
    unless @dry_run
      raise "command failed: '#{command}'" unless system command
    end
  end

  def sudo(command)
    sh "sudo #{command}"
  end

end

class BoxCommand
  include SystemTools

  def boot(options = {})
    qemu_options = []

    qemu_disks = []

    dist_dir = ["dist", options[:name]].compact.join("/")

    case options[:root]
    when "iso"
      qemu_options << "-cdrom #{dist_dir}/iso"
      qemu_options << "--boot order=d"
    else
      qemu_disks << "#{dist_dir}/disk"
    end

    qemu_disks.push *Dir["#{dist_dir}/storage*"]

    qemu_disks.each_with_index do |disk, index|
      qemu_options << "-drive file=#{disk},if=ide,index=#{index+1},media=disk"
    end

    mac_address_index = (options[:mac_index] or 0)
    mac_address = "52:54:00:12:34:#{56+mac_address_index}"
    
    qemu_options << "-net nic,macaddr=#{mac_address} -net vde,sock=/var/run/vde2/tap0.ctl"

    memory = (ENV['MEMORY'] or "512")

    ENV['QEMU_AUDIO_DRV']='alsa'

    qemu_command = "qemu -enable-kvm -m #{memory}m -soundhw ac97 #{qemu_options.join(' ')}"

    puts "Run #{qemu_command}"
    system qemu_command
  end

  class Cloner
    include SystemTools

    attr_accessor :target, :image, :dry_run, :skip_format, :external_extlinux

    def initialize(options = {}) 
      options = { :target => "/dev/sdb", :image => "dist/disk" }.merge options
      options.each { |k,v| send "#{k}=", v }
    end

    def partition
      "#{target}1"
    end

    def mounts
      IO.readlines("/proc/mounts").map { |m| m.scan(/\S+/) }
    end

    def mount(volume, options = {}, &block)
      name = File.basename(volume)

      unless target = options.delete(:target)
        Dir.mktmpdir("#{name}-") do |target|
          mount volume, {:target => target}.merge(options), &block
        end
        return
      end

      name = File.basename(volume)
      mount_options = options.collect do |pair|
        [ pair - [true] ].join('=')
      end.join(',')

      mount_options = "-o #{mount_options}" unless mount_options.empty?

      begin
        sudo "mount #{mount_options} #{volume} #{target}"
        yield target
      ensure
        sudo "umount #{target}"
      end
    end

    def mount_image
      mount(image, :loop => true, :offset => 64*512) do |mount_dir|
        yield mount_dir
      end
    end

    def chroot(root, command)
      sudo "mount proc #{root}/proc -t proc"
      sudo "mount -o bind /dev #{root}/dev"

      begin
        sudo "chroot #{root} #{command}"
      ensure
        sudo "umount #{root}/proc"
        sudo "umount #{root}/dev"
      end
    end

    def clone
      puts "Dry run mode" if dry_run
      confirm
      format unless skip_format
      copy
      extlinux
    end

    def extlinux
      mount(partition) do |partition_mount|
        if external_extlinux
          sudo "extlinux --install #{partition_mount}"
        else
          mount("#{partition_mount}/filesystem.squashfs", :loop => true) do |root|
            mount(partition_mount, :bind => true, :target => "#{root}/boot") do
              chroot root, "extlinux --install /boot"
            end
          end
        end
      end

      # Required ?
      # sudo "dd if=/usr/lib/syslinux/mbr.bin of=#{target}"
    end

    def copy
      puts "Copy files"
      mount_image do |image_mount|
        mount(partition) do |partition_mount|
          sudo "rsync -av #{image_mount}/ #{partition_mount}/"
        end
      end
    end

    def format
      puts "Formatting filesystem"
      sh "echo ',,L,*' | sudo /sbin/sfdisk -f -uS #{target}"
      sh "grep -q #{partition} /proc/mounts && sudo umount #{partition} || true"
      sudo "mke2fs -j -L boot #{partition}"
    end

    def confirm
      $stdout.write "Confirm you want install box image (#{image}) in #{target} [y/N] :"
      $stdout.flush
      exit 1 unless $stdin.read(1).downcase == 'y'
    end

  end

  def clone(options = {})
    Cloner.new(options).clone
  end

  def upgrade(options = {})
    host = options[:host]
    raise "No specified host for upgrade, use --host" unless host

    sh "scp dist/latest.yml dist/upgrade.tar #{host}:/tmp/"
    sh "ssh #{host} box-upgrade /tmp/upgrade.tar /tmp/latest.yml"
  end

end

BoxCommand.new.send(@command, @options)