# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""CounterManager for Mac OSX"""
from __future__ import absolute_import, print_function

import subprocess
import sys

from cmanager_base import CounterManager


def GetProcessData(pid):
    """Runs a ps on the process identified by pid and returns the output line
      as a list (pid, vsz, rss)
    """
    command = ['ps -o pid,vsize,rss -p'+str(pid)]
    try:
        handle = subprocess.Popen(command, stdout=subprocess.PIPE,
                                  universal_newlines=True, shell=True)
        handle.wait()
        data = handle.stdout.readlines()
    except Exception:
        print("Unexpected error executing '%s': %s", (command, sys.exc_info()))
        raise

    # First line is header output should look like:
    # PID      VSZ    RSS
    # 3210    75964    920
    line = data[1]
    line = line.split()
    if line[0] == str(pid):
        return line


def GetPrivateBytes(pid):
    """Calculate the amount of private, writeable memory allocated to a
    process.
    """
    psData = GetProcessData(pid)
    return int(psData[1]) * 1024  # convert to bytes


def GetResidentSize(pid):
    """Retrieve the current resident memory for a given process"""
    psData = GetProcessData(pid)
    return int(psData[2]) * 1024  # convert to bytes


class MacCounterManager(CounterManager):
    """This class manages the monitoring of a process with any number of
       counters.

       A counter can be any function that takes an argument of one pid and
       returns a piece of data about that process.
       Some examples are: CalcCPUTime, GetResidentSize, and GetPrivateBytes
    """

    counterDict = {"Private Bytes": GetPrivateBytes,
                   "RSS": GetResidentSize}

    def __init__(self, process_name, process, counters):
        """Args:
             counters: A list of counters to monitor. Any counters whose name
             does not match a key in 'counterDict' will be ignored.
        """

        CounterManager.__init__(self)

        # the last process is the useful one
        self.pid = process.pid

        self._loadCounters()
        self.registerCounters(counters)

    def getCounterValue(self, counterName):
        """Returns the last value of the counter 'counterName'"""
        if counterName not in self.registeredCounters:
            print("Warning: attempting to collect counter %s and it is not"
                  " registered" % counterName)
            return

        try:
            return self.registeredCounters[counterName][0](self.pid)
        except Exception as e:
            print("Error in collecting counter: %s, pid: %s, exception: %s"
                  % (counterName, self.pid, e))