require 'pty'
require 'expect'
require 'semantic_logger'

# Sftp that works on both Ruby and JRuby.
#
# Example Usage:
#   Sftp.upload(
#     user:            'friend',
#     password:        'secure',
#     host:            'localhost',
#     input_file_name: 'sample.txt'
#   )
#
# PTY.spawn is used so that the password can be passed via stdin into the sftp command line process.
#
# The sftp ruby gem was extremely slow and could not upload files larger than 2GB.
# This approach is much faster and does not suffer from the CPU and memory utilization issues
# associated with the sftp ruby gem.
#
# Note: The retry mechanism is only required on JRuby.
module Sftp
  include SemanticLogger::Loggable

  class SftpError < RuntimeError
  end

  class ConnectionFailure < SftpError
  end

  class SpawnFailure < SftpError
  end

  class ReadTimeout < SftpError
  end

  class SendFailure < SftpError
  end

  def self.upload(user:, password:, host:, port: 22, input_file_name:, output_file_name: input_file_name, path: nil, sftp_bin: 'sftp', connect_timeout_s: 30, read_timeout_s: 5)
    retry_count = 0
    PTY.spawn(sftp_bin, "-P#{port}", "#{user}@#{host}") do |reader, writer, pid|
      begin
        logger.trace "Pid: #{pid}"
        str = logger.measure_info("Connected using sftp process:[#{pid}] to #{host}:#{port}") do
          logger.trace 'Waiting for password'
          raise(SpawnFailure, "Failed to spawn #{sftp_bin}") unless reader.expect(/Password:.*/, read_timeout_s)
          writer.puts(password)
          logger.trace 'Waiting for Connected'
          raise(ConnectionFailure, 'Failed to connect') unless reader.expect(/Connected to #{host}.*/, connect_timeout_s)
          writer.puts "mkdir #{path}" if path

          writer.puts 'progress'
          logger.trace 'Waiting for Progress disabled'
          raise(ReadTimeout, 'Failed to disable progress') unless reader.expect(/Progress meter disabled.*/, read_timeout_s)
          writer.puts "put #{input_file_name} #{output_file_name}"
          logger.trace 'Waiting for put echo'
          raise(ReadTimeout, 'Failed to put file') unless reader.expect(/put.*/, read_timeout_s)
          logger.trace 'Waiting for blank line'
          logger.trace reader.gets # Skip remainder of put line above
          while (str = reader.gets).nil?
            sleep 0.1
            # Add limit
            logger.trace 'Waiting for more blank lines'
          end
          str
        end
        if str && str.include?('Uploading')
          logger.measure_info("Uploaded #{input_file_name} as #{output_file_name}") do
            writer.puts 'progress'
            logger.trace 'Waiting for Progress enabled (download complete)'
            # No timeout on the time it takes to send the file
            reader.expect(/Progress meter enabled.*/)
          end
        else
          raise(SendFailure, "Failed to upload #{input_file_name} as #{output_file_name}: #{str}")
        end
      ensure
        logger.measure_info('Waiting for sftp process to stop', min_duration: 5) do
          logger.trace 'puts bye'
          writer.puts 'bye'
          logger.trace 'writer close'
          writer.close
          logger.trace 'reader close'
          reader.close
          logger.trace 'process wait'
          Process.wait(pid, Process::WNOHANG)
          logger.trace "Finished. Status: #{$?.inspect}"
        end
      end
      true
    end
  rescue SpawnFailure
    retry_count += 1
    if retry_count <= 5
      logger.info "Retry attempt: #{retry_count}"
      retry
    end
  end

end