#!/usr/bin/env ruby

require 'rubygems'
require 'mysql'
require 'highline'
require 'optparse'

class CopyDatabase

  def initialize(options = {})
    @connection = Mysql.connect(options[:host], options[:username], options[:password])
    @connection.query_with_result = false
    @options = options
    Kernel.at_exit { @connection.close }
  end

  def copy(from_database, to_database)
    @from_database, @to_database, @tables = [from_database, to_database, nil]
    @connection.select_db(from_database)
    return if tables.empty?
    drop_database if @options[:drop_database]
    create_database if @options[:create_database]
    copy_structure
    lock_tables { copy_data }
  end

private

  def create_database
    query('CREATE DATABASE ' + "`#{@to_database}`")
  end

  def drop_database
    query('DROP DATABASE IF EXISTS ' + "`#{@to_database}`")
  end

  def copy_structure
    tables.each do |table|
      query('DROP TABLE IF EXISTS ' + sanitized_table(@to_database, table))
      query('CREATE TABLE ' + sanitized_table(@to_database, table) +
            ' LIKE ' + sanitized_table(@from_database, table))
    end
  end

  def copy_data
    tables.each do |table|
      disable_keys(@to_database, table) do
        query('INSERT INTO ' + sanitized_table(@to_database, table) +
              ' SELECT * FROM ' + sanitized_table(@from_database, table))
      end
    end
  end

  def disable_keys(database, table)
    query('ALTER TABLE ' + sanitized_table(database, table) + ' DISABLE KEYS')
    yield
    query('ALTER TABLE ' + sanitized_table(database, table) + ' ENABLE KEYS')
  end

  def lock_tables
    query('LOCK TABLES ' + locks.join(', '))
    yield
    query('UNLOCK TABLES')
  end

  def locks
    tables.inject(Array.new) do |arr, table|
      arr << (sanitized_table(@to_database, table) + ' WRITE') <<
             (sanitized_table(@from_database, table) + ' READ')
    end
  end

  def tables
    @tables ||= @connection.list_tables
  end

  def sanitized_table(database, table)
    "`#{database}`.`#{table}`"
  end

  def query(sql)
    puts sql if @options[:verbose]
    @connection.query(sql)
  end
end

class CopyDatabaseController

  def self.run(args)
    new(args).run
  end

  def initialize(args)
    @ui = HighLine.new
    parse(args)
  end

  def run
    CopyDatabase.new(@options).
                 copy(@from_database, @to_database)
  end

private

  def parse(args)
    @options = {
      :user => ENV['USER'],
      :password => '',
      :verbose => false,
      :host => '127.0.0.1',
      :create_database => false,
      :drop_database => false
    }

    parser = OptionParser.new do |opts|
      opts.banner = "Usage: #{File.basename($0)} [options] from-database to-database"
      opts.on('-u', '--user',
              'User to use when connecting to MySQL.',
              "Default: #{@options[:user]}") do |value|
        @options[:user] = value
      end
      opts.on('-p', '--password',
              'Prompt for a password to use when connecting to MySQL.') do
        @options[:password] = @ui.ask('Enter password: ') { |q| q.echo = false }
      end
      opts.on('-h', '--host HOST',
              'Host or socket to connect to.',
              "Default: #{@options[:host]}") do |value|
        @options[:host] = value
      end
      opts.on('-v', '--verbose', 'Prints queries as they are executed.') do
        @options[:verbose] = true
      end
      opts.on('-C', '--create-database',
              'Create the destination database before copying.') do
        @options[:create_database] = true
      end
      opts.on('-f', '--force',
              'Drop and then create the destination database before copying.') do
        @options[:drop_database] = true
        @options[:create_database] = true
      end
      opts.on('--help', 'Display this help message.') do
        puts opts
        exit
      end
    end

    parser.parse!(args)

    if args.size != 2
      puts parser
      exit
    end

    @from_database, @to_database = args
  end
end

CopyDatabaseController.run(ARGV)