module Selenium::WebDriver::TakesScreenshot

def save_screenshot(png_path, full_page: false)
  extension = File.extname(png_path).downcase
  if extension != '.png'
    WebDriver.logger.warn 'name used for saved screenshot does not match file type. ' \
                          'It should end with .png extension',
                          id: :screenshot
  end
  File.open(png_path, 'wb') { |f| f << screenshot_as(:png, full_page: full_page) }
end

def screenshot_as(format, full_page: false)
  case format
  when :base64
    full_page ? full_screenshot : screenshot
  when :png
    screenshot_as(:base64, full_page: full_page).unpack1('m')
  else
    raise Error::UnsupportedOperationError, "unsupported format: #{format.inspect}"
  end
rescue NameError
  raise Error::UnsupportedOperationError, "Full Page Screenshots are not supported for #{inspect}"
end