
from . import rich


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

  .. code-block:: python

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

  :options:

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

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

    return \
        {
            'selection': '',
            'queued': '',
        }


def read(data=None):
    r'''
  Read ``squeue -o "%all"``.

  :options:

    **data** (``<str>``)
      For debugging: specify the output of ``squeue -o "%all"`` 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('squeue -o "%all"', shell=True).decode('utf-8')

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

    # get the field-names
    head = head.split('|')

    # convert name (bug in slurm)
    # Bug #4948, https://bugs.schedmd.com
    idx = [i for i, key in enumerate(head) if key.strip() == 'USER']
    if len(idx) > 1:
        head[idx[1]] = 'USER_ID'

    # convert to list of dictionaries
    # - initialize
    lines = []
    # - loop over lines
    for line in data:
        # -- initialize empty dictionary
        info = {}
        # -- fill dictionary
        for key, val in zip(head, line.split('|')):
            if len(key.strip()) > 0:
                info[key.strip()] = val.strip()
        # -- store to list of lines
        lines += [info]

    # return output
    return lines


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

  :arguments:

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

  :options:

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

  :returns:

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

    import time

    # read time
    if now is None:
        now = time.mktime(time.localtime())

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

    # loop over all lines
    for line in lines:

        # convert to integer
        for key in ['CPUS', 'NODES']:
            line[key] = rich.Integer(line[key])

        # "year-month-dayThour:minute:second" (e.g. "2017-11-05T19:09:53") -> seconds from now
        for key in ['START_TIME', 'SUBMIT_TIME']:
            try:
                line[key] = rich.Duration(
                    int(now - time.mktime(time.strptime(line[key], '%Y-%m-%dT%H:%M:%S'))))
            except BaseException:
                line[key] = rich.Duration(line[key])

        # "days-hours:mins:secs" (e.g. "1-4:18:13") -> seconds
        for key in ['TIME_LIMIT', 'TIME_LEFT', 'TIME']:
            line[key] = rich.Duration(line[key])

        # convert memory (e.g. "4G") -> bytes
        for key in ['MIN_MEMORY']:
            line[key] = rich.Memory(line[key])

        # specialize number of CPUS
        if str(line['ST']) == 'R':
            line['CPUS_R'] = rich.Integer(int(line['CPUS']))
            line['CPUS_PD'] = rich.Integer(0)
        else:
            line['CPUS_R'] = rich.Integer(0)
            line['CPUS_PD'] = rich.Integer(int(line['CPUS']))

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

        # highlight queued jobs
        if str(line['ST']) == 'PD':
            for key in line:
                line[key].color = theme['queued']

    return lines


def read_interpret(data=None, now=None, theme=colors()):
    r'''
  Read and interpret ``squeue -o "%all"``.

  :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), now, theme)
