#!/usr/bin/env python
# ----------------------------------------------------------------------
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ----------------------------------------------------------------------

import sys, re, os, logging, configparser
from optparse import OptionParser, OptionGroup


# ----------------------------------------------------------------------
# Constants & variables
# ----------------------------------------------------------------------

RAINBOW_VERSION = '2.5.0'
"""
The script version number.

@type: string
"""

RAINBOW_CONFIGS_HOME = os.path.join(os.sep,'usr','share','rainbow','configs')
"""
The location where rainbow configs are stored.

@type: string
"""

USER_CONFIGS_HOME = os.path.expanduser('~/.rainbow')
"""
The location where rainbow user-defined configs are stored.

@type: string
"""

USAGE = "%prog [options] -- command [args...] "
"""
A string showing the script usage syntax.

@type: string
"""

DESCRIPTION = """Colorize commands output using patterns. If you don't specify a command, STDIN is used.
For more information, see rainbow man page (man rainbow)."""
"""
A string describing the script.

@type: string
"""

FILTERS = {
  'red':                { 'short_option': 'r', 'help': 'display RED pattern in red',                                   'filter_start': chr(27)+'[31m', 'filter_end': chr(27)+'[39m' },
  'green':              { 'short_option': 'g', 'help': 'display GREEN pattern in green',                               'filter_start': chr(27)+'[32m', 'filter_end': chr(27)+'[39m' },
  'yellow':             { 'short_option': 'y', 'help': 'display YELLOW pattern in yellow',                             'filter_start': chr(27)+'[33m', 'filter_end': chr(27)+'[39m' },
  'blue':               { 'short_option': 'b', 'help': 'display BLUE pattern in blue',                                 'filter_start': chr(27)+'[34m', 'filter_end': chr(27)+'[39m' },
  'magenta':            { 'short_option': 'm', 'help': 'display MAGENTA pattern in magenta',                           'filter_start': chr(27)+'[35m', 'filter_end': chr(27)+'[39m' },
  'cyan':               { 'short_option': 'c', 'help': 'display CYAN pattern in cyan',                                 'filter_start': chr(27)+'[36m', 'filter_end': chr(27)+'[39m' },
  'start-red':          {                      'help': 'toggle foreground color to red at START_RED pattern',          'filter_start': chr(27)+'[31m', 'filter_end': ''             },
  'start-green':        {                      'help': 'toggle foreground color to green at START_GREEN pattern',      'filter_start': chr(27)+'[32m', 'filter_end': ''             },
  'start-yellow':       {                      'help': 'toggle foreground color to yellow at START_YELLOW pattern',    'filter_start': chr(27)+'[33m', 'filter_end': ''             },
  'start-blue':         {                      'help': 'toggle foreground color to blue at START_BLUE pattern',        'filter_start': chr(27)+'[34m', 'filter_end': ''             },
  'start-magenta':      {                      'help': 'toggle foreground color to magenta at START_MAGENTA pattern',  'filter_start': chr(27)+'[35m', 'filter_end': ''             },
  'start-cyan':         {                      'help': 'toggle foreground color to cyan at START_CYAN pattern',        'filter_start': chr(27)+'[36m', 'filter_end': ''             },
  'reset-color':        {                      'help': 'reset foreground color at RESET_COLOR pattern',                'filter_start': ''            , 'filter_end': chr(27)+'[39m' },
  'bred':               { 'short_option': 'R', 'help': 'display BRED pattern in red background',                       'filter_start': chr(27)+'[41m', 'filter_end': chr(27)+'[49m' },
  'bgreen':             { 'short_option': 'G', 'help': 'display BGREEN pattern in green background',                   'filter_start': chr(27)+'[42m', 'filter_end': chr(27)+'[49m' },
  'byellow':            { 'short_option': 'Y', 'help': 'display BYELLOW pattern in yellow background',                 'filter_start': chr(27)+'[43m', 'filter_end': chr(27)+'[49m' },
  'bblue':              { 'short_option': 'B', 'help': 'display BBLUE pattern in blue background',                     'filter_start': chr(27)+'[44m', 'filter_end': chr(27)+'[49m' },
  'bmagenta':           { 'short_option': 'M', 'help': 'display BMAGENTA pattern in magenta background',               'filter_start': chr(27)+'[45m', 'filter_end': chr(27)+'[49m' },
  'bcyan':              { 'short_option': 'C', 'help': 'display BCYAN pattern in cyan background',                     'filter_start': chr(27)+'[46m', 'filter_end': chr(27)+'[49m' },
  'start-bcyan':        {                      'help': 'toggle background color to cyan at START_BCYAN pattern',       'filter_start': chr(27)+'[46m', 'filter_end': ''             },
  'start-bred':         {                      'help': 'toggle background color to red at START_BRED pattern',         'filter_start': chr(27)+'[41m', 'filter_end': ''             },
  'start-bgreen':       {                      'help': 'toggle background color to green at START_BGREEN pattern',     'filter_start': chr(27)+'[42m', 'filter_end': ''             },
  'start-byellow':      {                      'help': 'toggle background color to yellow at START_BYELLOW pattern',   'filter_start': chr(27)+'[43m', 'filter_end': ''             },
  'start-bblue':        {                      'help': 'toggle background color to blue at START_BBLUE pattern',       'filter_start': chr(27)+'[44m', 'filter_end': ''             },
  'start-bmagenta':     {                      'help': 'toggle background color to magenta at START_BMAGENTA pattern', 'filter_start': chr(27)+'[45m', 'filter_end': ''             },
  'reset-bcolor':       {                      'help': 'reset background color at RESET_BCOLOR pattern',               'filter_start': ''            , 'filter_end': chr(27)+'[49m' },
  'bold':               {                      'help': 'display BOLD pattern in bold',                                 'filter_start': chr(27)+'[1m',  'filter_end': chr(27)+'[22m' },
  'faint':              {                      'help': 'display FAINT pattern with decreased intensity',               'filter_start': chr(27)+'[2m',  'filter_end': chr(27)+'[22m' },
  'start-bold':         {                      'help': 'toggle bold on at START_BOLD pattern',                         'filter_start': chr(27)+'[1m',  'filter_end': ''             },
  'start-faint':        {                      'help': 'toggle faint on at START_FAINT pattern',                       'filter_start': chr(27)+'[2m',  'filter_end': ''             },
  'reset-intensity':    {                      'help': 'reset text intensity at RESET_INTENSITY pattern',              'filter_start': ''           ,  'filter_end': chr(27)+'[22m' },
  'italic':             {                      'help': 'display ITALIC pattern in italic',                             'filter_start': chr(27)+'[3m',  'filter_end': chr(27)+'[23m' },
  'underline':          {                      'help': 'display UNDERLINE pattern underlined',                         'filter_start': chr(27)+'[4m',  'filter_end': chr(27)+'[24m' },
  'underline-double':   {                      'help': 'display UNDERLINE_DOUBLE pattern double underlined',           'filter_start': chr(27)+'[21m', 'filter_end': chr(27)+'[24m' },
  'blink':              {                      'help': 'display BLINK pattern blinking',                               'filter_start': chr(27)+'[5m',  'filter_end': chr(27)+'[25m' },
  'blink-rapid':        {                      'help': 'display BLINK pattern blinking',                               'filter_start': chr(27)+'[6m',  'filter_end': chr(27)+'[25m' },
  'negative':           {                      'help': 'display NEGATIVE pattern swapping foreground and background',  'filter_start': chr(27)+'[7m',  'filter_end': chr(27)+'[27m' },
  'hide':               {                      'help': 'hide HIDE pattern',                                            'filter_start': chr(27)+'[8m',  'filter_end': chr(27)+'[28m' }
 }
"""
A dictionary of the filters available in the scripts.

The filter name defines the command line long option name. The other fields are self-explanatory.

@type: dict
"""

FILTER_GROUPS = [
  { 'name': 'Foreground color',
    'help': 'Use these options to associate patterns to text foreground colors.',
    'filters': [ 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'start-red', 'start-green', 'start-yellow', 'start-blue', 'start-magenta', 'start-cyan', 'reset-color' ] },
  { 'name': 'Background color',
    'help': 'Use these options to associate patterns to text background colors.',
    'filters': [ 'bred', 'bgreen', 'byellow', 'bblue', 'bmagenta', 'bcyan', 'start-bred', 'start-bgreen', 'start-byellow', 'start-bblue', 'start-bmagenta', 'start-bcyan', 'reset-bcolor' ] },
  { 'name': 'Text intensity',
    'help': 'Use these options to associate patterns to text emphasis.',
    'filters': [ 'bold', 'faint', 'start-bold', 'start-faint', 'reset-intensity' ] },
  { 'name': 'Other formattings',
    'help': 'Use these options to associate patterns to text formattings. Some of the filters may not work on all terminals.',
    'filters': [ 'italic', 'underline', 'underline-double', 'blink', 'blink-rapid', 'negative', 'hide' ] }
 ]
"""
A dictionary defining groups of filters.

This is used to make the C{--help} option output easier to read.

@type: string
"""

logger = logging.getLogger("RAINBOW")
"""
Script logger.

@type: Logger
"""

optionParser = OptionParser(usage=USAGE, version="%prog " + RAINBOW_VERSION, description=DESCRIPTION)
"""
The parser used to handle command line arguments.

@type: OptionParser
"""

patterns = {}
"""
The dictionary containing patterns registered through configs and command line options.

@type: dict
"""


# ----------------------------------------------------------------------
# Functions
# ----------------------------------------------------------------------


def register_pattern_with_filter(pattern,filter_name):
    """
    Register a pattern to be processed by a filter.

    This functions populates the L{patterns} dictionary.

    @param pattern: the pattern (regular expression).
    @type pattern: string
    @param filter_name: the filter to associate with the pattern. It must be a key of L{FILTERS}.
    @type filter_name: string
    @rtype: void
    """

    logger.debug("Binding pattern '%s' with filter '%s'.", pattern, filter_name)

    # If this is new pattern, create a new entry with the corresponding compiled regex.
    if pattern not in patterns:
        patterns[pattern] = { 'pattern_start': '',
                              'pattern_end': '',
                              'regex' : re.compile(pattern) }

    # Update the pattern before and after strings with the filter ones.
    patterns[pattern]['pattern_start'] += FILTERS[filter_name]['filter_start']
    patterns[pattern]['pattern_end'] += FILTERS[filter_name]['filter_end']


def handle_command_line_pattern_option(option, opt, value, parser):
    """
    Handle a command line option defining a pattern/filter association.
    
    @param option: the Option instance calling the callback.
    @type option: Option
    @param opt: the option string seen on the command-line.
    @type opt: string
    @param value: the argument to this option seen on the command-line.
    @type value: string
    @param parser: the OptionParser instance.
    @type parser: OptionParser
    @rtype: void
    """

    # The pattern is the option value, the filter name is computed from the option name ("--filter").
    register_pattern_with_filter(value,option.get_opt_string()[2:])


def handle_command_line_verbosity_option(option, opt, value, parser):
    """
    Handle a command line option increasing the logger verbosity.

    @param option: the Option instance calling the callback.
    @type option: Option
    @param opt: the option string seen on the command-line.
    @type opt: string
    @param value: the argument to this option seen on the command-line.
    @type value: string
    @param parser: the OptionParser instance.
    @type parser: OptionParser
    @rtype: void
    """

    # Decrease the logger level.
    logger.setLevel(logger.level - 10)

def locate_config(config):
    """
    Try to locate a config and return the corresponding file path.

    @param config: the config file to look for. The function will look for (in the following order):
        - C{config}
        - C{config.cfg}
        - C{L{USER_CONFIGS_HOME}/config}
        - C{L{USER_CONFIGS_HOME}/config.cfg}
        - C{L{RAINBOW_CONFIGS_HOME}/config}
        - C{L{RAINBOW_CONFIGS_HOME}/config.cfg}
    @type config: string
    @rtype: string
    """

    logger.debug("Trying to locate config '%s'.", config)

    # Try to locate the config.
    for dir in [os.path.curdir, USER_CONFIGS_HOME, RAINBOW_CONFIGS_HOME]:
        if os.path.isfile(os.path.join(dir,config)):
            return  os.path.join(dir,config)
        elif os.path.isfile(os.path.join(dir,config + ".cfg")):
            return  os.path.join(dir,config + ".cfg")

def process_config(file,options):
    """
    Load a config file and register the imports, patterns and options it defines.

    @param file: the config file to look for. The function will look for (in the following order):
        - C{config}
        - C{config.cfg}
        - C{L{USER_CONFIGS_HOME}/config}
        - C{L{USER_CONFIGS_HOME}/config.cfg}
        - C{L{RAINBOW_CONFIGS_HOME}/config}
        - C{L{RAINBOW_CONFIGS_HOME}/config.cfg}
    @type file: string
    @param options
    @type options: string
    @rtype: void
    """

    logger.debug("Parsing the config '%s'.", file)

    # Parse the configuration file
    configParser = configparser.ConfigParser()
    configParser.read(file)

    # Process general section
    if configParser.has_section('general'):

        # Process imports
        if configParser.has_option('general', 'imports'):
            for config_import in [v.strip() for v in configParser.get('general', 'imports').split(',')]:
                config_file = locate_config(config_import)
                if config_file:
                    process_config(config_file,options)
                else:
                    logger.error("Could not locate the config '%s'.", config_import)

        # Process options
        if configParser.has_option('general', 'enable-stderr-filtering'):
            options.enable_stderr_filtering = configParser.get('general', 'enable-stderr-filtering')

    # Process filters section
    if configParser.has_section('filters'):
        for filter in configParser.options("filters"):
            if filter in FILTERS:
                register_pattern_with_filter(configParser.get("filters",filter),filter)
            else:
                logger.warning("Unknown filter '%s'.", filter)

    logger.info("Loaded config '%s'.", file)


def apply_filters(line):
    """
    Colorize the line using the L{patterns} dictionary.

    @param line: the line to process.
    @type line: string
    @return: the processed line.
    @rtype: string
    """

    # Look for each pattern registered.
    for (pattern,pattern_filter) in list(patterns.items()):

        # Apply filter for each match.
        for match in pattern_filter['regex'].finditer(line):
            line = line.replace(match.group(),pattern_filter['pattern_start'] + match.group() + pattern_filter['pattern_end'])

    return line.rstrip()


def main():
    """
    Rainbow main program.

    @return: the exit code of the program.
    @rtype: int
    """

    # Setup the logger.
    logger_console_handler = logging.StreamHandler()
    logger_formatter = logging.Formatter("[%(name)s|%(levelname)s] %(message)s")
    logger_console_handler.setFormatter(logger_formatter)
    logger.addHandler(logger_console_handler)
    logger.setLevel(logging.WARNING)

    # Setup the command line option parser.
    optionParser.formatter.max_help_position = 50
    optionParser.formatter.width = 150

    # Register the command line options.
    optionParser.add_option("-f",
                            "--config",
                            action="append",
                            dest="config",
                            type="string",
                            help="Load a config file defining patterns. Go to %s for examples. The option can be called several times." % (RAINBOW_CONFIGS_HOME))
    optionParser.add_option("-v",
                            "--verbose",
                            action="callback",
                            callback=handle_command_line_verbosity_option,
                            help="Turn on verbose mode. This option can be called several times to increase the verbosity level.")
    optionParser.add_option("--disable-stderr-filtering",
                            action="store_false", dest="enable_stderr_filtering",
                            default=True,
                            help="Disable STDERR filtering, which can have unexpected effects on commands directly using tty.")

    # Walk through filter groups and register each filter as a command line option.
    for filter_group in FILTER_GROUPS:
        filter_option_group = OptionGroup(optionParser, filter_group['name'], filter_group['help'])
        for filter_name in filter_group['filters']:
            filter = FILTERS[filter_name]
            if 'short_option' in filter:
                filter_option_group.add_option("-" + filter['short_option'],
                                               "--" + filter_name, action="callback",
                                               callback=handle_command_line_pattern_option,
                                               type="string",
                                               help=filter['help'])
            else:
                filter_option_group.add_option("--" + filter_name,
                                               action="callback",
                                               callback=handle_command_line_pattern_option,
                                               type="string",
                                               help=filter['help'])
        optionParser.add_option_group(filter_option_group)

    # Parse command line options.
    (options, args) = optionParser.parse_args()

    # If configs passed, process them.
    if options.config:
        for config in options.config:
            config_file = locate_config(config)
            if config_file:
                logger.debug("Config '%s' located at '%s'.", config, config_file)
                process_config(config_file,options)
            else:
                logger.error("Could not locate the config '%s'.", config)

    try:

        if args:

            # If no pattern defined after parsing options and configs, lookup a default config from the target command name.
            if len(patterns) == 0:
                logger.debug("No pattern defined, looking for a default config for '%s'." % args[0])
                config_file = locate_config(args[0])
                if config_file:
                    logger.info("No pattern defined, the config '%s' will be loaded." % config_file)
                    process_config(config_file,options)
                else:
                    logger.debug("Could not locate a default config for '%s'.", args[0])

            # Create a pipe.
            pipe_read, pipe_write = os.pipe()

            # Fork rainbow execution in a child process and execute the target command in the parent process.
            pid = os.fork()

            # In the child process, connect the pipe and execute the target command.
            if pid == 0:
                os.close(pipe_read)
                os.dup2(pipe_write, sys.stdout.fileno())
                if options.enable_stderr_filtering == True:
                    os.dup2(pipe_write, sys.stderr.fileno())
                os.execvp(args[0],args)

            # In the parent process, apply filters on each line of the pipe output.
            else:
                os.close(pipe_write)
                pipe_fd_handle = os.fdopen(pipe_read)
                while True:
                    try:
                        line = pipe_fd_handle.readline()
                        if line:
                            print(apply_filters(line[:-1]))
                        else:
                            break
                    except KeyboardInterrupt:
                        pass
                os._exit(os.EX_OK)

        else:
            logger.info("No arguments given, using STDIN as input.")
            try:
                while True:
                    print(apply_filters(input()))
            except (EOFError, KeyboardInterrupt):
                return 0

    # Raise an error for any other exception.
    except Exception as exc:
        logger.error("%s", exc)
        return 1


# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------

if __name__ == "__main__":
    sys.exit(main())
