require 'json'
require 'open3'

def vector_add(a, b, device: 'cpu', python: ENV.fetch('PYTHON', 'python3'))
  worker = File.join(__dir__, 'worker.py')
  output, error, status = Open3.capture3(
    python, worker, '--device', device,
    stdin_data: JSON.generate({a: a, b: b})
  )
  raise "worker failed: #{error.strip}" unless status.success?

  result = JSON.parse(output)
  unless result.is_a?(Hash) && result['values'].is_a?(Array) &&
         result['device'] == device && result['dtype'] == 'float32'
    raise 'invalid worker response'
  end
  result
end

if $PROGRAM_NAME == __FILE__
  puts JSON.generate(vector_add([1, -2, 0.5], [3, 2, 0.25]))
end