"""
Profiling class for updated information on process status
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

__RCSID__ = "$Id$"

import datetime
import errno
import psutil

from DIRAC import gLogger, S_OK, S_ERROR
from DIRAC.Core.Utilities.DErrno import EEZOMBIE, EENOPID, EEEXCEPTION


def checkInvocation(func):
  """ Decorator for invoking psutil methods
  """
  def wrapper(*args, **kwargs):
    try:
      return func(*args, **kwargs)
    except psutil.ZombieProcess as e:
      gLogger.error('Zombie process: %s' % e)
      return S_ERROR(EEZOMBIE, 'Zombie process: %s' % e)
    except psutil.NoSuchProcess as e:
      gLogger.error('No such process: %s' % e)
      return S_ERROR(errno.ESRCH, 'No such process: %s' % e)
    except psutil.AccessDenied as e:
      gLogger.error('Access denied: %s' % e)
      return S_ERROR(errno.EPERM, 'Access denied: %s' % e)
    except Exception as e:  # pylint: disable=broad-except
      gLogger.error(e)
      return S_ERROR(EEEXCEPTION, e)

  return wrapper


class Profiler(object):
  """
  Class for profiling both general stats about a machine and individual processes.
  Every instance of this class is associated to a single process by using its PID.
  Calls to the different methods of the class will return the current state of the process.
  """

  def __init__(self, pid=None):
    """
    :param str pid: PID of the process to be profiled
    """
    self.process = None
    if pid:
      try:
        self.process = psutil.Process(int(pid))
      except psutil.NoSuchProcess as e:
        gLogger.error('No such process: %s' % e)

  def pid(self):
    """
    Returns the process PID
    """
    if self.process:
      return S_OK(self.process.pid)
    else:
      gLogger.error('No PID of process to profile')
      return S_ERROR(EENOPID, 'No PID of process to profile')

  @checkInvocation
  def status(self):
    """ Returns the process status
    """
    result = self.process.status()
    return S_OK(result)

  @checkInvocation
  def runningTime(self):
    """
    Returns the uptime of the process
    """
    start = datetime.datetime.fromtimestamp(self.process.create_time())
    result = (datetime.datetime.now() - start).total_seconds()
    return S_OK(result)

  @checkInvocation
  def memoryUsage(self, withChildren=False):
    """
    Returns the memory usage of the process in MB
    """
    # Information is returned in bytes
    rss = self.process.memory_info().rss
    if withChildren:
      for child in self.process.children(recursive=True):
        rss += child.memory_info().rss
    # converted to MB
    return S_OK(rss / float(2 ** 20))

  @checkInvocation
  def vSizeUsage(self, withChildren=False):
    """
    Returns the memory usage of the process in MB
    """
    # Information is returned in bytes
    vms = self.process.memory_info().vms
    if withChildren:
      for child in self.process.children(recursive=True):
        vms += child.memory_info().vms
    # converted to MB
    return S_OK(vms / float(2 ** 20))

  @checkInvocation
  def numThreads(self, withChildren=False):
    """
    Returns the number of threads the process is using
    """
    nThreads = self.process.num_threads()
    if withChildren:
      for child in self.process.children(recursive=True):
        nThreads += child.num_threads()
    return S_OK(nThreads)

  @checkInvocation
  def cpuPercentage(self, withChildren=False):
    """
    Returns the percentage of cpu used by the process
    """
    cpuPercentage = self.process.cpu_percent()
    if withChildren:
      for child in self.process.children(recursive=True):
        cpuPercentage += child.cpu_percent()
    return S_OK(cpuPercentage)

  @checkInvocation
  def cpuUsageUser(self, withChildren=False, withTerminatedChildren=False):
    """
    Returns the percentage of cpu used by the process
    """
    cpuUsageUser = self.process.cpu_times().user
    childrenUser = 0
    oldChildrenUser = 0
    if withChildren:  # active children
      for child in self.process.children(recursive=True):
        childrenUser += child.cpu_times().user
      gLogger.debug(
          "CPU user (process, children)", "(%.1fs, %.1fs)" %
          (cpuUsageUser, childrenUser))
    if withTerminatedChildren:  # all terminated children of the root process
      for child in self.process.children(recursive=True):
        oldChildrenUser += child.cpu_times().children_user
      gLogger.debug(
          "CPU user (process, old children)", "(%.1fs, %.1fs)" %
          (cpuUsageUser, oldChildrenUser))
    else:
      gLogger.debug("CPU user", "%.1fs" % cpuUsageUser)
    return S_OK(cpuUsageUser + childrenUser + oldChildrenUser)

  @checkInvocation
  def cpuUsageSystem(self, withChildren=False, withTerminatedChildren=False):
    """
    Returns the percentage of cpu used by the process
    """
    cpuUsageSystem = self.process.cpu_times().system
    childrenSystem = 0
    oldChildrenSystem = 0
    if withChildren:  # active children
      for child in self.process.children(recursive=True):
        childrenSystem += child.cpu_times().system
      gLogger.debug(
          "CPU user (process, children)", "(%.1fs, %.1fs)" %
          (cpuUsageSystem, childrenSystem))
    if withTerminatedChildren:  # all terminated children of the root process
      for child in self.process.children(recursive=True):
        oldChildrenSystem += child.cpu_times().children_system
      gLogger.debug(
          "CPU user (process, old children)", "(%.1fs, %.1fs)" %
          (cpuUsageSystem, oldChildrenSystem))
    else:
      gLogger.debug("CPU user", "%.1fs" % cpuUsageSystem)
    return S_OK(cpuUsageSystem + childrenSystem + oldChildrenSystem)

  def getAllProcessData(self, withChildren=False, withTerminatedChildren=False):
    """
    Returns data available about a process
    """
    data = {}

    data['datetime'] = datetime.datetime.utcnow()
    data['stats'] = {}

    result = self.pid()
    if result['OK']:
      data['stats']['pid'] = result['Value']

    result = self.status()
    if result['OK']:
      data['stats']['status'] = result['Value']

    result = self.runningTime()
    if result['OK']:
      data['stats']['runningTime'] = result['Value']

    result = self.memoryUsage(withChildren)
    if result['OK']:
      data['stats']['memoryUsage'] = result['Value']

    result = self.vSizeUsage(withChildren)
    if result['OK']:
      data['stats']['vSizeUsage'] = result['Value']

    result = self.numThreads(withChildren)
    if result['OK']:
      data['stats']['threads'] = result['Value']

    result = self.cpuPercentage(withChildren)
    if result['OK']:
      data['stats']['cpuPercentage'] = result['Value']

    result = self.cpuUsageUser(withChildren, withTerminatedChildren)
    if result['OK']:
      data['stats']['cpuUsageUser'] = result['Value']

    result = self.cpuUsageSystem(withChildren, withTerminatedChildren)
    if result['OK']:
      data['stats']['cpuUsageSystem'] = result['Value']

    return S_OK(data)
