#!/usr/bin/env python3
# Copyright 2019 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This script computes the number of concurrent jobs that can run in the
# build as a function of the machine. It accepts a set of key value pairs
# given by repeated --memory-per-job arguments. For example:
#
# $ get_concurrent_jobs.py --memory-per-job dart=1GB
#
# The result is a json map printed to stdout that gives the number of
# concurrent jobs allowed of each kind. For example:
#
# {"dart": 8}
#
# Some memory can be held out of the calculation with the --reserve-memory flag.

import argparse
import ctypes
import json
import multiprocessing
import os
import re
import subprocess
import sys

UNITS = {'B': 1, 'KB': 2**10, 'MB': 2**20, 'GB': 2**30, 'TB': 2**40}


# pylint: disable=line-too-long
# See https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-globalmemorystatusex
# and https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/ns-sysinfoapi-memorystatusex
# pylint: enable=line-too-long
class MEMORYSTATUSEX(ctypes.Structure):
  _fields_ = [
      ('dwLength', ctypes.c_ulong),
      ('dwMemoryLoad', ctypes.c_ulong),
      ('ullTotalPhys', ctypes.c_ulonglong),
      ('ullAvailPhys', ctypes.c_ulonglong),
      ('ullTotalPageFile', ctypes.c_ulonglong),
      ('ullAvailPageFile', ctypes.c_ulonglong),
      ('ullTotalVirtual', ctypes.c_ulonglong),
      ('ullAvailVirtual', ctypes.c_ulonglong),
      ('sullAvailExtendedVirtual', ctypes.c_ulonglong),
  ]


def get_total_memory():
  if sys.platform in ('win32', 'cygwin'):
    stat = MEMORYSTATUSEX(dwLength=ctypes.sizeof(MEMORYSTATUSEX))
    success = ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(stat))
    return stat.ullTotalPhys if success else 0
  if sys.platform.startswith('linux'):
    if os.path.exists('/proc/meminfo'):
      with open('/proc/meminfo') as meminfo:
        memtotal_re = re.compile(r'^MemTotal:\s*(\d*)\s*kB')
        for line in meminfo:
          match = memtotal_re.match(line)
          if match:
            return float(match.group(1)) * 2**10
  if sys.platform == 'darwin':
    try:
      return int(subprocess.check_output(['sysctl', '-n', 'hw.memsize']))
    except:  # pylint: disable=bare-except
      return 0
  return 0


def parse_size(string):
  i = next(i for (i, c) in enumerate(string) if not c.isdigit())
  number = string[:i].strip()
  unit = string[i:].strip()
  return int(float(number) * UNITS[unit])


class ParseSizeAction(argparse.Action):

  def __call__(self, parser, args, values, option_string=None):
    sizes = getattr(args, self.dest, [])
    for value in values:
      (k, val) = value.split('=', 1)
      sizes.append((k, parse_size(val)))
    setattr(args, self.dest, sizes)


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--memory-per-job',
      action=ParseSizeAction,
      default=[],
      nargs='*',
      help='Key value pairings (dart=1GB) giving an estimate of the amount of '
      'memory needed for the class of job.'
  )
  parser.add_argument(
      '--reserve-memory',
      type=parse_size,
      default=0,
      help='The amount of memory to be held out of the amount for jobs to use.'
  )
  args = parser.parse_args()

  total_memory = get_total_memory()

  # Ensure the total memory used in the calculation below is at least 0
  mem_total_bytes = max(0, total_memory - args.reserve_memory)

  # Ensure the number of cpus used in the calculation below is at least 1
  try:
    cpu_cap = multiprocessing.cpu_count()
  except:  # pylint: disable=bare-except
    cpu_cap = 1

  concurrent_jobs = {}
  for job, memory_per_job in args.memory_per_job:
    # Calculate the number of jobs that will fit in memory. Ensure the
    # value is at least 1.
    num_concurrent_jobs = int(max(1, mem_total_bytes / memory_per_job))
    # Cap the number of jobs by the number of cpus available.
    concurrent_jobs[job] = min(num_concurrent_jobs, cpu_cap)

  print(json.dumps(concurrent_jobs))

  return 0


if __name__ == '__main__':
  sys.exit(main())