#!/usr/bin/env ruby

require 'correct_horse_battery_staple'
require 'correct_horse_battery_staple/assembler'
require 'commander/import'
require 'logger'

program :version, "0.1.0"
program :description, "Correct Horse Battery Staple - XKCD-style Passphrases"

global_option '-D', '--debug', 'Display full protocol trace'
global_option '-V', '--verbose', 'Shows request and response'

logger = Logger.new(STDERR)

default_command :generate

command :repl do |c|
  c.syntax = 'repl'
  c.summary = 'Open Pry-based REPL, optionally with a loaded corpus'
  c.option '-c CORPUSFILE', '--corpus CORPUSFILE', "File to use as the source corpus"
  c.action do |args, options|
    require 'pry'

    if options.corpus
      format = CorrectHorseBatteryStaple::Corpus.format_for(options.corpus)
      corpus = CorrectHorseBatteryStaple.load_corpus(options.corpus, format)
    end
    
    binding.pry
  end
end

command :convert do |c|
  c.syntax = 'convert <infile> [<outfile>]'
  c.summary = 'Convert a corpus from one format to another'
  c.option '-i IFORMAT', '--iformat IFORMAT', 'Set the output corpus format'
  c.option '-o OFORMAT', '--oformat OFORMAT', 'Set the output corpus format'
  c.option '-r', '--recalculate', 'Recalculate all statistics before outputting'
  c.option '-S FIELD', '--sortby FIELD', 'Sort words by field'
  c.option '-R', '--randomize', 'Randomize word order'
  c.option '-l COUNT', '--limit COUNT', Integer, 'Set the maximum words processed'
  c.option '-P', '--precache', 'Precache the source word list in memory'

  c.action do |args, options|
    options.default :limit => -1

    infile = args[0]
    outfile = args[1] && args[1] != "-" ?
              (!args[1].include?(':') ? open(args[1], "w") : args[1]) :
              STDOUT

    options.iformat ||= CorrectHorseBatteryStaple::Corpus.format_for(infile)
    unless options.oformat
      if args[1]
        options.oformat = CorrectHorseBatteryStaple::Corpus.format_for(args[1])
      else
        raise ArgumentError, "You must supply an output format via the file extension or -o option"
      end
    end

    corpus = CorrectHorseBatteryStaple.load_corpus(args[0], options.iformat)
    if options.precache
      corpus.precache(16_000_000)
    end

    # CSV doesn't currently persist stats
    corpus.recalculate if options.recalculate || options.iformat.to_s == "csv"

    corpus.table.shuffle! if options.randomize
    if options.limit > -1
      corpus.table.slice!(options.limit..-1) if corpus.size > options.limit
    end

    with_tempfile(outfile) do |tempfile|
      writer = CorrectHorseBatteryStaple::Writer.make_writer(tempfile, options.oformat)
      writer.write_corpus(corpus)
      writer.close
    end

  end
end

command :stats do |c|
  c.syntax = 'stats <corpusfile>'
  c.summary = 'Show statistics from a corpus'

  c.action do |args,options|
    corpus = CorrectHorseBatteryStaple::Corpus.read(args[0])

    stats = corpus.stats
    stats.to_hash.each do |key, val|
      puts "#{key}: #{val}"
    end
  end
end

command :inspect do |c|
  c.syntax = 'inspect <corpusfile>'
  c.summary = 'Show implementation information for a corpus'

  c.action do |args,options|
    corpus = CorrectHorseBatteryStaple::Corpus.read(args[0])
    puts corpus.inspect
  end
end

command :generate do |c|
  c.syntax = 'generate [<number of words>]'
  c.summary = 'Generate a passphrase with a given number of words'
  #c.description = ''
  #c.example 'description', 'command example'
  c.option '-f FORMAT', '--format FORMAT', 'Set the corpus format'
  c.option '-c CORPUSFILE', '--corpus CORPUSFILE', "File to use as the source corpus"
  c.option '-n COUNT', '--repeat COUNT', Integer, 'How many passwords to generate'
  c.option '-P', '--precache', 'Precache the word list in memory'

  # filters
  c.option '-W MIN..MAX', '--wordsize MIN..MAX',
           'Set the allowed word size in number of characters'
  c.option '-P MIN..MAX', '--percentile MIN..MAX',
           'Set the percentile range of the word frequency'

  c.action do |args, options|
    options.default :wordsize => "3..7", :repeat => 1,
                    :corpus => CorrectHorseBatteryStaple::DEFAULT_CORPUS_NAME,
                    :percentile => "30..80"

    number_of_words = (args[0] || 4).to_i
    count = options.repeat.to_i

    corpus = CorrectHorseBatteryStaple.load_corpus(options.corpus, options.format)
    if options.precache
      corpus.precache(16_000_000)
    end
    puts "loaded corpus!" if options.debug
    unfiltered_size = corpus.count

    make_options = {}

    f_wordsize = CorrectHorseBatteryStaple::RangeParser.new.parse options.wordsize
    unless f_wordsize.begin <= 1 && f_wordsize.end >= 30
      make_options[:word_length] = f_wordsize
    end

    f_percentile = CorrectHorseBatteryStaple::RangeParser.new.parse options.percentile
    unless f_percentile.include? 0..100
      make_options[:percentile] = f_percentile
    end

    if options.verbose
      entropy = number_of_words * corpus.entropy_per_word.floor
      guesses_per_sec = 1000.0
      search_space = 2**entropy
      years = search_space / guesses_per_sec / (365*24*3600)
      puts "Corpus size: #{corpus.length} candidate words of #{unfiltered_size} total"
      puts "Entropy: #{entropy} bits (2^#{entropy} = #{search_space})"
      puts "Years to guess at #{guesses_per_sec.to_i} guesses/sec: #{years.round}"
    end

    generator = CorrectHorseBatteryStaple::Generator.new(corpus)

    count.times do
      puts generator.make(number_of_words, make_options)
    end
  end
end

command :list do |c|
  c.syntax = 'list'
  c.summary = 'Show a list of available corpus names'
  c.option '-p', '--paths', 'Show paths'

  c.action do |args, options|
    list = CorrectHorseBatteryStaple.corpus_list(:with_paths => options.paths)

    list.each_with_index do |item, i|
      list[i] = "[#{item}]" if item.include? CorrectHorseBatteryStaple::DEFAULT_CORPUS_NAME
    end
    puts list.sort.join(options.paths ? "\n" : ', ')
  end
end

command :mkcorpus do |c|
  c.syntax = 'mkcorpus file [file, file, ...]'
  c.summary = 'Generate a JSON or CSV corpus from an input file'
  #c.description = ''
  #c.example 'description', 'command example'
  c.option '-o FILENAME', '--output FILENAME', 'Set the output filename or spec'
  c.option '-I FORMAT', '--iformat FORMAT', 'Set the input format - "wiktionary" or "wordfrequency"'
  c.option '-O FORMAT', '--oformat FORMAT', 'Set the output format - "marshal", "json", or "csv"'
  c.option '-R', '--randomize', 'Randomize word order'
  c.option '-l COUNT', '--limit COUNT', Integer, 'Set the maximum words processed'
  c.action do |args, options|
    options.default :iformat => "wiktionary", :limit => -1

    assembler = CorrectHorseBatteryStaple::Assembler.new(
      CorrectHorseBatteryStaple::Parser::Regex.new(options.iformat.to_sym))
    assembler.read(args)
    assembler.randomize if options.randomize
    if options.limit > -1
      assembler.limit(options.limit)
    end
    corpus = assembler.corpus

    destination = options.output || "-"

    options.oformat ||= CorrectHorseBatteryStaple::Corpus.format_for(destination)

    with_tempfile(destination) do |tempname|
      CorrectHorseBatteryStaple::Writer.write(corpus, tempname, options.oformat)
    end
  end
end

def with_tempfile(realname)
  pathpart = realname.respond_to?(:path) ? realname.path : realname
  tempname = nil
  if (realname.is_a?(String) ||
      (realname.is_a?(File) && File.exist?(realname.path))) &&
      ! pathpart != '-' &&
      ! pathpart.include?(':')
    begin
      tempname = File.join(File.dirname(pathpart), ".temp-#{File.basename(pathpart)}")
      yield tempname
      File.rename tempname, pathpart
    rescue
      File.delete(tempname) if tempname && File.exist?(tempname)
      raise
    end
  else
    yield realname
  end
end