#!/usr/bin/env ruby

require 'optparse'

@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("-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

class BoxCommand

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

    qemu_disks = []

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

    qemu_disks.push *Dir["dist/storage*"]

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

    qemu_options << "-net nic -net vde,sock=/var/run/vde2/tap0.ctl"

    ENV['QEMU_AUDIO_DRV']='alsa'

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

    puts "Run #{qemu_command}"
    system qemu_command
  end

  def clone(options = {})
    target = (options[:target] or "/dev/sdb")
    disk_image = (options[:image] or "dist/disk")
    partition = "#{target}1"
    tmpdir = "/tmp/dist"

    sh "[ -d #{tmpdir} ]||mkdir #{tmpdir}"
    if disk_image =~ /([a-z]+box)-(\d+-\d+)/
      # like pigebox-20101014-1155
      name = $1
      release = $2
      disk_image = "#{tmpdir}/#{name}-#{release}.disk"
      sh "ssh dev.tryphon.priv cat /var/lib/buildbot/dist/#{name}/#{name}-#{release}.disk.gz | gunzip -c > #{disk_image}" unless File.exist?(disk_image)
    end

    raise "Can't find #{disk_image}" unless File.exists?(disk_image)

    #if partition_mount = mounts.assoc(partition)
    #  partition_mount_point = partition_mount[1]
    #else
      partition_mount_point = "#{tmpdir}/boot"
      #sudo "mkdir #{partition_mount_point}" unless File.exist? partition_mount_point
      #sudo "mount #{partition} #{partition_mount_point}"
    #end

    #unless File.exists? "#{partition_mount_point}/config.pp"
      $stdout.write "Confirm you want install box image (#{disk_image}) in #{target} [y/N] :"
      $stdout.flush
      exit 1 unless $stdin.read(1).downcase == 'y'
    #end

    if mounts.assoc(partition)
      puts "Unmount #{target}"
      exit 1 unless system "sudo umount #{partition}"
    end

    unless options[:skip_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 -L boot #{partition}"
    end

    sh "[ -d #{partition_mount_point} ]||mkdir #{partition_mount_point}"
    
    puts "Copy files"
    sudo "mount #{partition} #{partition_mount_point}"
    mount_disk_image_fs(disk_image) do |dir|
      sudo "rsync -av #{dir}/ #{partition_mount_point}/"
    end
    sudo "umount #{partition_mount_point}"

    puts "Install extlinux"
    sudo "mount #{partition} #{partition_mount_point}"
    #sudo "mount proc #{partition_mount_point}/proc -t proc"
    #sudo "mount -o bind /dev #{partition_mount_point}/dev"

    begin
      #sudo "chroot #{tmpdir} /usr/bin/extlinux --install /boot"
      sudo "/usr/sbin/extlinux --install #{partition_mount_point}"
    ensure
      #sudo "umount #{tmpdir}/proc"
      #sudo "umount #{tmpdir}/dev"
      sudo "umount #{partition_mount_point}"
    end

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

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

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

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

  def mount_disk_image_fs(disk_image)
    mount_dir = "/tmp/mount_boot_fs"
    sh "[ -d #{mount_dir} ]||mkdir -p #{mount_dir}"
    begin
      sudo "mount -o loop,offset=#{64*512} #{disk_image} #{mount_dir}"
      yield mount_dir
    ensure
      sudo "umount #{mount_dir}"
    end
  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)