
from . import rich


def colors(theme=None):
    r'''
Return dictionary of colors.

.. code-block:: python

    {
      'selection' : '...',
    }

:options:

    **theme** ([``'dark'``] | ``<str>``)
        Select color-theme.
    '''

    if theme == 'dark':
        return \
            {
                'selection': '1;32;40',
            }

    return \
        {
            'selection': '',
        }


def read(data=None):
    r'''
Read ``ps -eo pid,user,rss,%cpu,command``.

:options:

    **data** (``<str>``)
        For debugging: specify the output of ``ps -eo pid,user,rss,%cpu,command`` as string.

:returns:

    **lines** ``<list<dict>>``
        A list of dictionaries, that contain the different fields. All data are strings.
    '''

    import subprocess

    # get live info
    if data is None:
        data = subprocess.check_output(
            'ps -eo pid,user,rss,%cpu,command',
            shell=True).decode('utf-8')

    # extract the header and the info
    header, data = data.split('\n', 1)
    data = list(filter(None, data  .split('\n')))
    header = list(filter(None, header.split(' ')))

    # convert to list of dictionaries
    # - initialize
    lines = []
    # - loop over lines
    for line in data:
        # -- decode line with information
        d = ['' for i in header]
        line = line.strip()
        d[0], line = line.split(' ', 1)
        line = line.strip()
        d[1], line = line.split(' ', 1)
        line = line.strip()
        d[2], line = line.split(' ', 1)
        line = line.strip()
        d[3], line = line.split(' ', 1)
        d[4] = line.strip()
        # -- initialize empty dictionary
        info = {}
        # -- fill dictionary
        for key, val in zip(header, d):
            info[key] = val
        # -- store to list of lines
        lines += [info]

    # return output
    return lines


def interpret(lines, theme=colors()):
    r'''
Interpret the output of ``GooseSLURM.ps.read``. All fields are converted to the
``GooseSLURM.rich`` classes adding useful colors in the process.

:arguments:

    **lines** ``<list<dict>>``
        The output of ``GooseSLURM.ps.read``

:options:

    **theme** (``<dict>``)
        The color-theme, as selected by ``GooseSLURM.ps.colors``.

:returns:

    **lines** (``<list<dict>>``)
        A list of dictionaries, that contain the different fields. All data are
        ``GooseSLURM.rich.String`` or derived types.
    '''

    # interpret input as list
    if not isinstance(lines, list):
        lines = [lines]

    # loop over all lines
    for line in lines:

        # custom conversion
        for key in ['%CPU']:
            line[key] = rich.Float(line[key], precision=2)

        # custom conversion
        for key in ['RSS']:
            line[key] = rich.Memory(line[key], default_unit=1e3)

        # convert remaining fields to string
        for key in line:
            if not isinstance(line[key], rich.String):
                line[key] = rich.String(line[key])

    return lines


def read_interpret(data=None, theme=colors()):
    r'''
Read and interpret ``ps -eo pid,user,rss,%cpu,command``.

:returns:

    **lines** (``<list<dict>>``)
        A list of dictionaries, that contain the different fields. All data are
        ``GooseSLURM.rich.String`` or derived types.
    '''

    return interpret(read(data), theme)
