
from . import rich

import re


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

.. code-block:: python

    {
        'selection' : '...',
        'free'      : '...',
        'error'     : '...',
        'warning'   : '...',
        'low'       : '...',
    }

:options:

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

    if theme == 'dark':
        return \
            {
                'selection': '1;32;40',
                'free': '1;32',
                'error': '9;31',
                'warning': '1;31',
                'low': '1;36',
            }

    return \
        {
            'selection': '',
            'free': '',
            'error': '',
            'warning': '',
            'low': '',
        }


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

:options:

    **data** (``<str>``)
        For debugging: specify the output of ``sinfo -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('sinfo -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 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 cpu_score(CPU_LOAD, CPUS_A, **kwargs):

    if not CPU_LOAD.isnumeric():
        return rich.Float('')
    if int(CPUS_A) <= 0:
        return rich.Float('')

    return rich.Float(float(CPU_LOAD) / float(CPUS_A), precision=2)


def mem_score(MEMORY, FREE_MEM, CPUS_A, CPUS_T, **kwargs):

    if not FREE_MEM.isnumeric():
        return rich.Float('')
    if not CPUS_T  .isnumeric():
        return rich.Float('')
    if int(MEMORY) <= 0:
        return rich.Float('')
    if int(CPUS_A) <= 0:
        return rich.Float('')

    MEM_USED = float(MEMORY) - float(FREE_MEM)

    return rich.Float(float(MEM_USED) / float(MEMORY) * float(CPUS_T) / float(CPUS_A), precision=2)


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

:arguments:

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

:options:

    **theme** (``<dict>``)
        The color-theme, as selected by ``GooseSLURM.sinfo.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:

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

                # covert to float
        for key in ['CPU_LOAD']:
            line[key] = rich.Float(str(line[key]))

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

        # convert memory (e.g. "4G") -> bytes
        for key in ['MEMORY', 'FREE_MEM']:
            line[key] = rich.Memory(str(line[key]), default_unit=1e6)

        # CPUs
        # - split: allocated/idle/other/total
        a, i, o, t = str(line['CPUS(A/I/O/T)']).split('/')
        # - store extra fields
        line['CPUS_A'] = rich.Integer(a)
        line['CPUS_I'] = rich.Integer(i)
        line['CPUS_O'] = rich.Integer(o)
        line['CPUS_T'] = rich.Integer(t)

        # compute scores
        line['CPU_RELJOB'] = cpu_score(**line)
        line['MEM_RELJOB'] = mem_score(**line)

        # default: CPUs down/online
        line['CPUS_D'] = rich.Integer(0)
        line['CPUS_O'] = rich.Integer(str(line['CPUS_T']))

        # node down: mark all fields
        if re.match(r'down,*', str(line['STATE'])) or \
           re.match(r'maint.*', str(line['STATE'])) or \
           re.match(r'drain.*', str(line['STATE'])):

            line['CPUS_I'] = rich.Integer(0)
            line['CPUS_D'] = rich.Integer(str(line['CPUS_T']))
            line['CPUS_O'] = rich.Integer(0)

            for key in line:
                line[key].color = theme['error']

        # highlight 'scores'
        if int(line['CPUS_I']) > 0:
            line['CPUS_I'].color = theme['free']
        if float(line['CPU_RELJOB']) > 1.05:
            line['CPU_RELJOB'].color = theme['warning']
        elif float(line['CPU_RELJOB']) < 0.95:
            line['CPU_RELJOB'].color = theme['low']

        # highlight 'scores'
        if float(line['MEM_RELJOB']) > 0.9:
            line['MEM_RELJOB'].color = theme['warning']

    return lines


def read_interpret(data=None, theme=colors()):
    r'''
Read and interpret ``sinfo -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), theme)
