require 'find'
require 'rubygems'
require 'digest/md5'
gem 'net-sftp', '<2.0.0'
require 'net/sftp'

class SftpDirPublisher
  CHECKSUM_FILENAME = ".checksums"

  attr_reader :host, :username, :password, :remote_dir, :local_dir, :exclude

  def initialize(host, username, password, remote_dir, local_dir, exclude=[])
    @host = host
    @username = username
    @password = password
    @remote_dir = remote_dir
    @local_dir = local_dir
    @verbose = true
    @exclude = exclude
  end

  def upload(dry_run=false)
    create_checksums
    puts_if_verbose "Connecting to #{@host}..."
    Net::SFTP.start(@host, @username, @password, :timeout=>30) do |sftp|

      puts_if_verbose "Fetching checksums..."
      local_checksums = YAML::load(File.open( local_path(CHECKSUM_FILENAME) ))
      remote_checksums = {}
      begin
        checksums_src = ''
        sftp.open_handle( server_path(CHECKSUM_FILENAME) ) do |handle|
          checksums_src = sftp.read(handle)
        end
        remote_checksums = YAML::load(checksums_src) unless checksums_src.nil? or checksums_src==''
      rescue Net::SFTP::Operations::StatusException=>se
        # server's checksum.yml is missing
      end

      puts_if_verbose "Comparing checksum data..."
      to_upload, to_remove = hash_diff(local_checksums, remote_checksums)

      if to_upload.length > 0 or to_remove.length > 0
        puts_if_verbose "Differences found:"
        to_upload.each {|f| puts_if_verbose " - (Upload) #{f}" }
        to_remove.each {|f| puts_if_verbose " - (Delete) #{f}"}

        exit if dry_run

        puts_if_verbose "Beginning sync..."

        to_upload.each do |filename|
          begin
            puts_if_verbose " - #{remote_path(filename)}"

            dir_name = File.dirname(filename)
            dir_segs = dir_name.split('/')
            #puts "Checking path: #{dir_segs.join( '/' )}"
            prog_path = []
            dir_segs.each do |partial_dir|
              begin
                prog_path << partial_dir
                sftp.mkdir( remote_path( prog_path ), :permissions=>0755 )
                puts_if_verbose " + #{remote_path( prog_path )}"
              rescue Net::SFTP::Operations::StatusException=>se
                # don't worry about it
              end
            end

            sftp.put_file local_path(filename), remote_path(filename)
            sftp.open_handle( remote_path(filename) ) do |handle|
              sftp.fsetstat( handle, :permissions=>0644 )
            end
          rescue  Net::SFTP::Operations::StatusException=>se
            puts_if_verbose " ! Error uploading '#{filename}': #{se}"
            puts_if_verbose; puts_if_verbose "Halted execution of upload."
            exit(1)
          end
        end

        to_remove.each do |filename|
          begin
            sftp.remove remote_path(filename)
            puts_if_verbose " x #{remote_path(filename)}"
          rescue
            puts_if_verbose " ! Error removing '#{filename}': #{$!}"
          end
        end

        begin
          sftp.put_file local_path(CHECKSUM_FILENAME), remote_path(CHECKSUM_FILENAME)
          sftp.open_handle( remote_path(CHECKSUM_FILENAME) ) do |handle|
            sftp.fsetstat( handle, :permissions=>0644 )
          end
        rescue
          puts_if_verbose " ! Error uploading '#{CHECKSUM_FILENAME}': #{$!}"
        end


        summary = "#{to_upload.length} file(s) uploaded"
        summary += ", #{to_remove.length} files(s) deleted" if to_remove.length > 0

        puts summary
      else
        puts "No changes made. The server is up to date!"
      end

      puts "Done."
    end
  end

protected

  def create_checksums
    checksums = generate_from(@local_dir)
    File.open(local_path(CHECKSUM_FILENAME), 'w') do |f|
      f.write checksums.to_yaml
    end
    checksums
  end

  # returns [files_to_upload, files_to_delete]
  def hash_diff(source={}, target={})
    # look for differences...
    src_files = source.fetch('files', {})
    tgt_files = target.fetch('files', {})
    to_update = []; to_delete = []; to_upload = []

    tgt_files.each do |filename, checksum|
      if src_files.has_key? filename
        to_update << filename unless src_files[filename] == checksum
      else
        to_delete << filename
      end
    end

    to_upload = src_files.keys - tgt_files.keys

    # returns [files_to_upload, files_to_delete]
    [[to_upload, to_update].flatten.sort, to_delete.sort]
  end

  # Create the checksums...
  def generate_from(cache_dir)
    checksums = { 'generated_on'=>Time.now, 'files'=>{} }
    checksum_path = local_path(CHECKSUM_FILENAME)
    puts "Generating checksums from #{cache_dir}"
    Find.find( cache_dir ) do |f|
      next if File.directory?( f ) or f == checksum_path or has_exclusion(f,cache_dir)

      checksums['files'][f.gsub("#{cache_dir}/", '')] = ::Digest::MD5.hexdigest( File.read(f) )
    end
    checksums
  end

  def has_exclusion(path, cache_dir)
    path = path.gsub("#{cache_dir}/", '')
    exclude.any? do |test_pattern|
      case test_pattern
        when String: path.start_with?(test_pattern)
        when Regexp: path =~ test_pattern
      end
    end
  end

  def server_path(path)
    [@remote_dir, path].flatten.join('/')
  end
  alias_method :remote_path, :server_path

  def local_path(path)
    File.join(@local_dir, path)
  end

  def puts_if_verbose(msg='')
    puts(msg) if @verbose
  end

end

class String
  def starts_with?(prefix)
    prefix = prefix.to_s
    self[0, prefix.length] == prefix
  end
end