from __future__ import print_function

from argparse import ArgumentParser
import os
import sys
import threading

_color = threading.local()
_color.enabled = True


__all__ = ['CustomRegistryMeta', 'HelpfulArgumentParser', 'disable_color', 'color_enabled', 'colored', 'cprint', 'STATUS_COLORS']


STATUS_COLORS = {
    'new': 'grey',
    'pending': 'grey',
    'running': 'yellow',
    'successful': 'green',
    'failed': 'red',
    'error': 'red',
    'canceled': 'grey',
}


class CustomRegistryMeta(type):
    @property
    def registry(cls):
        return dict((command.name, command) for command in cls.__subclasses__())


class HelpfulArgumentParser(ArgumentParser):
    def error(self, message):  # pragma: nocover
        """Prints a usage message incorporating the message to stderr and
        exits.
        If you override this in a subclass, it should not return -- it
        should either exit or raise an exception.
        """
        self.print_help(sys.stderr)
        self._print_message('\n')
        self.exit(2, '%s: %s\n' % (self.prog, message))


def color_enabled():
    return _color.enabled


def disable_color():
    _color.enabled = False


COLORS = dict(
    list(
        zip(
            [
                'grey',
                'red',
                'green',
                'yellow',
                'blue',
                'magenta',
                'cyan',
                'white',
            ],
            list(range(30, 38)),
        )
    )
)


def colored(text, color=None):
    '''Colorize text w/ ANSI color sequences'''
    if _color.enabled and os.getenv('ANSI_COLORS_DISABLED') is None:
        fmt_str = '\033[%dm%s'
        if color is not None:
            text = fmt_str % (COLORS[color], text)
        text += '\033[0m'
    return text


def cprint(text, color, **kwargs):
    if _color.enabled:
        print(colored(text, color), **kwargs)
    else:
        print(text, **kwargs)
