#!/usr/bin/env ruby
$: << File.expand_path("#{File.dirname __FILE__}/../lib")
require 'rubygems'
require 'dyna'
require 'optparse'

Version = Dyna::VERSION

mode = nil
file = 'Dynafile'
output_file = '-'
split = false
MAGIC_COMMENT = <<-EOS
# -*- mode: ruby -*-
# vi: set ft=ruby :
EOS

options = {
  :dry_run => false,
  :format  => :ruby,
  :color   => true,
  :debug   => false,
}

ARGV.options do |opt|
  begin
    access_key = nil
    secret_key = nil
    region = nil
    profile_name = nil
    credentials_path = nil
    format_passed = false
    role_arn = nil
    serial_number = nil
    token_code = nil
    role_session_name = nil

    opt.on('-p', '--profile PROFILE_NAME')                   {|v| profile_name                  = v      }
    opt.on('',   '--role_arn ROLE_ARN')                      {|v| role_arn                      = v      }
    opt.on('',   '--serial_number SERIAL_NUMBER')            {|v| serial_number                 = v      }
    opt.on('',   '--token_code TOKEN_CODE')                  {|v| token_code                    = v      }
    opt.on('',   '--role_session_name ROLE_SESSION_NAME')    {|v| role_session_name             = v      }
    opt.on(''  , '--credentials-path PATH')                  {|v| credentials_path              = v      }
    opt.on('-k', '--access-key ACCESS_KEY')                  {|v| access_key                    = v      }
    opt.on('-s', '--secret-key SECRET_KEY')                  {|v| secret_key                    = v      }
    opt.on('-r', '--region REGION')                          {|v| region                        = v      }
    opt.on('-l', '--use-local')                              {    options[:use_local]           = true   }
    opt.on('',   '--endpoint DYNAMO_ENDPOINT')               {|v| options[:dynamo_endpoint]     = v      }
    opt.on('-a', '--apply')                                  {    mode                          = :apply }
    opt.on('-f', '--file FILE')                              {|v| file                          = v      }
    opt.on('',   '--table_names TABLE_NAMES', Array)         {|v| options[:table_names]         = v      }
    opt.on('',   '--exclude_table_names TABLE_NAMES', Array) {|v| options[:exclude_table_names] = v      }
    opt.on('',   '--dry-run')                                {    options[:dry_run]             = true   }
    opt.on('-e', '--export')                                 {    mode                          = :export}
    opt.on('-o', '--output FILE')                            {|v| output_file                   = v      }
    opt.on('',   '--split')                                  {    split                         = true   }
    opt.on('',   '--split-more')                             {    split                         = :more  }
    opt.on(''  , '--no-color')                               {    options[:color]               = false  }
    opt.on(''  , '--debug')                                  {    options[:debug]               = true   }
    opt.parse!

    aws_opts = {}
    if options[:use_local]
      options[:access_key_id] = access_key if access_key
      options[:secret_access_key] = secret_key if secret_key
      options[:region] = region if region
    elsif access_key and secret_key
      aws_opts = {
        :access_key_id     => access_key,
        :secret_access_key => secret_key,
      }
    elsif profile_name or credentials_path
      credentials_opts = {}
      credentials_opts[:profile_name] = profile_name if profile_name
      credentials_opts[:path] = credentials_path if credentials_path
      provider = AWS::Core::CredentialProviders::SharedCredentialFileProvider.new(credentials_opts)
      aws_opts[:credential_provider] = provider
    elsif role_arn and serial_number
      credentials = Aws::AssumeRoleCredentials.new(
        client: Aws::STS::Client.new,
        role_arn: role_arn,
        role_session_name: role_session_name,
        serial_number: serial_number,
        token_code: token_code,
      )
      aws_opts[:credentials] = credentials
    elsif (access_key and !secret_key) or (!access_key and secret_key) or mode.nil?
      puts opt.help
      exit 1
    end

    aws_opts[:region] = region if region
    Aws.config.update(aws_opts)

    # Remap groups to exclude to regular expressions (if they're surrounded by '/')
    if options[:exclude_table_names]
      options[:exclude_table_names].map! do |name|
        name =~ /\A\/(.*)\/\z/ ? Regexp.new($1) : Regexp.new("\A#{Regexp.escape(name)}\z")
      end
    end
  rescue => e
    $stderr.puts("[ERROR] #{e.message}")
    exit 1
  end
end

String.colorize = options[:color]

if options[:debug]
  Aws.config.update({
    :http_wire_trace => true,
    :logger => Dyna::Logger.instance,
  })
end

begin
  logger = Dyna::Logger.instance
  logger.set_debug(options[:debug])
  client = Dyna::Client.new(options)

  case mode
  when :export
    if split
      logger.info('Export Table')

      output_file = 'Dynafile' if output_file == '-'
      requires = []
      base_dir = File.dirname(output_file)

      client.export(options) do |exported, converter|
        write_table_file = proc do |table_file, tables|
          requires << table_file

          logger.info("  write `#{table_file}`")
          FileUtils.mkdir_p(File.dirname(table_file))

          open(table_file, 'wb') do |f|
            f.puts MAGIC_COMMENT
            f.puts converter.call(tables)
          end
        end

        exported.each do |table_name, tables|
          table_file = File.join(base_dir, "#{Aws::DynamoDB::Client.new.config.region}/#{table_name}.dyna")
          write_table_file.call(table_file, table_name => tables)
        end
      end

      logger.info("  write `#{output_file}`")
      open(output_file, 'wb') do |f|
        f.puts MAGIC_COMMENT

        requires.each do |table_file|
          table_file.sub!(%r|\A#{Regexp.escape base_dir}/?|, '')
          f.puts "require '#{table_file}'"
        end
      end
    else
      exported = client.export(options)

      if output_file == '-'
        logger.info('# Export Table') if options[:format] == :ruby
        puts exported
      else
        logger.info("Export Table to `#{output_file}`")

        open(output_file, 'wb') do |f|
          f.puts MAGIC_COMMENT if options[:format] == :ruby
          f.puts exported
        end
      end
    end
  when :apply
    unless File.exist?(file)
      raise "No Dynafile found (looking for: #{file})"
    end

    msg = "Apply `#{file}` to DynamoDB"
    msg << ' (dry-run)' if options[:dry_run]
    logger.info(msg)

    updated = client.apply(file)

    logger.info('No change'.intense_blue) unless updated
  else
    raise 'must not happen'
  end
rescue => e
  if options[:debug]
    raise e
  else
    $stderr.puts("[ERROR] #{e.message}".red)
    $stderr.puts("#{e.backtrace.join("\n")}".red)
    exit 1
  end
end