#!/usr/bin/env ruby
require 'json'
require 'socket'

class SendReceive
  COMMANDS = %w[skel base incremental transfer cleanup cancel].freeze

  def self.run(argv: ARGV, env: ENV)
    if argv.length != 2
      warn 'Usage: $0 <pool> <key name>'
      exit(false)
    end

    new(
      argv[0],
      argv[1],
      env.fetch('SSH_ORIGINAL_COMMAND', nil),
      env['SSH_CONNECTION'].split
    )
  end

  def initialize(pool, key_name, cmdline, connection)
    @key_pool = pool
    @key_name = key_name
    @client_ip, = connection
    error! unless cmdline

    args = cmdline.split
    error! if args.count < 3 || args[0] != 'receive'

    @protocol_version = Integer(args[1], exception: false)
    error! unless @protocol_version

    command = args[2]
    error! unless COMMANDS.include?(command)

    @args = args[3..]
    connect
    method(command.to_sym).call
  end

  protected

  attr_reader :key_pool, :key_name, :client_ip, :protocol_version, :args, :client

  def skel
    to_pool =
      case args[0]
      when nil, '-'
        nil
      else
        args[0]
      end

    send_cmd(
      :receive_skel,
      pool: to_pool,
      passphrase: args[1],
      client_ip:,
      key_pool:,
      key_name:
    )

    if recv_resp! != 'continue'
      warn 'Error: invalid response'
      exit(false)
    end

    send_stdin

    # osctld will generate a unique token for this send/receive, which we pass
    # to the sending side for further identification
    puts recv_resp!
  end

  def base
    error! if args.count < 2
    send_cmd(
      :receive_base,
      key_pool:,
      key_name:,
      token: parse_token,
      dataset: args[1],
      snapshot: args[2]
    )

    if recv_resp! != 'continue'
      warn 'Error: invalid response'
      exit(false)
    end

    send_stdin
    recv_resp!
  end

  def incremental
    error! if args.count < 2
    send_cmd(
      :receive_incremental,
      key_pool:,
      key_name:,
      token: parse_token,
      dataset: args[1],
      snapshot: args[2]
    )

    if recv_resp! != 'continue'
      warn 'Error: invalid response'
      exit(false)
    end

    send_stdin
    recv_resp!
  end

  def transfer
    send_cmd(
      :receive_transfer,
      key_pool:,
      key_name:,
      token: parse_token,
      start: args[1] == 'start'
    )
    recv_resp!
  end

  def cleanup
    send_cmd(
      :receive_cleanup,
      key_pool:,
      key_name:,
      token: parse_token
    )
    recv_resp!
  end

  def cancel
    send_cmd(
      :receive_cancel,
      key_pool:,
      key_name:,
      token: parse_token
    )
    recv_resp!
  end

  def connect
    @client = UNIXSocket.new('/run/osctl/send-receive/control.sock')
  end

  def send_cmd(cmd, opts = {})
    client.puts({ cmd:, opts: opts.merge(protocol_version:) }.to_json)
  end

  def send_stdin
    client.send_io($stdin)
  end

  def recv_msg
    JSON.parse(client.readline, symbolize_names: true)
  end

  def recv_resp!
    loop do
      msg = recv_msg

      unless msg[:status]
        warn "Error: #{msg[:message]}"
        exit(false)
      end

      next if msg.has_key?(:progress)

      return msg[:response]
    end
  end

  def parse_token
    error! unless args[0]
    args[0]
  end

  def usage
    warn <<~END
      Usage:
        receive <protocol> skel [pool|- [passphrase]]
        receive <protocol> base <token> <dataset> [snapshot]
        receive <protocol> incremental <token> <dataset> [snapshot]
        receive <protocol> transfer <token> [start]
        receive <protocol> cleanup <token>
        receive <protocol> cancel <token>
    END
  end

  def error!
    usage
    exit(false)
  end
end

SendReceive.run if $0 == __FILE__
