# Copyright 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import atexit
import inspect
import logging
import os

from collections import defaultdict


def GetChildPids(processes, pid):
  """Returns all child processes of |pid| from the given |processes| list.

  Args:
    processes: A tuple of (pid, ppid, state) as generated by ps.
    pid: The pid for which to get children.

  Returns:
    A list of child pids.
  """
  child_dict = defaultdict(list)
  for curr_pid, curr_ppid, state in processes:
    if 'Z' in state:
      continue  # Ignore zombie processes
    child_dict[int(curr_ppid)].append(int(curr_pid))
  queue = [pid]
  child_ids = []
  while queue:
    parent = queue.pop()
    if parent in child_dict:
      children = child_dict[parent]
      queue.extend(children)
      child_ids.extend(children)
  return child_ids


def GetPsOutputWithPlatformBackend(platform_backend, columns, pid):
  """Returns output of the 'ps' command as a list of lines.

  Args:
    platform_backend: The platform backend (LinuxBasedPlatformBackend or
        PosixPlatformBackend).
    columns: A list of require columns, e.g., ['pid', 'pss'].
    pid: If not None, returns only the information of the process with the pid.
  """
  args = ['ps']
  args.extend(['-p', str(pid)] if pid != None else ['-e'])
  for c in columns:
    args.extend(['-o', c + '='])
  return platform_backend.RunCommand(args).splitlines()


def EnableListingStrayProcessesUponExitHook():
  def _ListAllSubprocesses():
    try:
      import psutil
    except ImportError:
      logging.warning(
          'psutil is not installed on the system. Not listing possible '
          'leaked processes. To install psutil, see: '
          'https://pypi.python.org/pypi/psutil')
      return
    telemetry_pid = os.getpid()
    parent = psutil.Process(telemetry_pid)
    if hasattr(parent, 'children'):
      children = parent.children(recursive=True)
    else:  # Some old version of psutil use get_children instead children.
      children = parent.get_children()
    if children:
      leak_processes_info = []
      for p in children:
        if inspect.ismethod(p.name):
          name = p.name()
        else:  # Process.name is a property in old versions of psutil.
          name = p.name
        process_info = '%s (%s)' % (name, p.pid)
        try:
          if inspect.ismethod(p.cmdline):
            cmdline = p.cmdline()
          else:
            cmdline = p.cmdline
          process_info += ' - %s' % cmdline
        except Exception as e:
          logging.warning(str(e))
        leak_processes_info.append(process_info)
      logging.error('Telemetry leaks these processes: %s',
                    ', '.join(leak_processes_info))

  atexit.register(_ListAllSubprocesses)