#!/usr/bin/env python

# Copyright 2017 The Fuchsia Authors
#
# Use of this source code is governed by a MIT-style
# license that can be found in the LICENSE file or at
# https://opensource.org/licenses/MIT

"""

This tool will color the lines from loglistener

Example usage #1:
  # Colorize log lines, with all messages from the same thread in the same color
  loglistener | scripts/colorize_logs -t

Example usage #2:
  # Colorize log lines, with all messages from the same process in the same color
  loglistener | scripts/colorize_logs

Example usage #3:
  # Print the colorization of log.txt to stdout
  # Identical to `scripts/colorize_logs < log.txt`
  scripts/colorize_logs log.txt

Example usage #4:
  # Colorize all ERROR and INFO lines in log.txt
  scripts/colorize_logs -r ERROR -r INFO log.txt

Example usage #5:
  # Colorize all lines with drv='<something>' in log.txt with distinct colors
  # for each <something>
  scripts/colorize_logs -r "drv='[^']*'" log.txt

"""

import argparse
import re
import sys

BASE_COLORS = [
    #'\033[40m', # black
    '\033[91m', # red
    '\033[92m', # green
    '\033[93m', # yellow
    '\033[94m', # blue
    '\033[95m', # magenta
    '\033[66m', # cyan
    #'\033[47m', # white
]
RESET_BG = '\033[49m'
RESET_FG = '\033[39m'

class ColorAssigner(object):
  def __init__(self, colors):
    self.lru = list(colors)
    self.task_colors = { }

  def get_bg_color(self, task_id):
    if task_id not in self.task_colors:
      c = self.lru.pop(0)
      self.task_colors[task_id] = c
    else:
      c = self.task_colors[task_id]
      self.lru.remove(c)

    self.lru.append(c)
    return c


PROCESS_RE = r'^\[\d+\.\d+] (\d+)[\.:]\d+> .*$'
THREAD_RE = r'^\[\d+\.\d+] (\d+[\.:]\d+)> .*$'

def main():
  parser = argparse.ArgumentParser(
      description=__doc__,
      formatter_class=argparse.RawDescriptionHelpFormatter)
  parser.add_argument("--process", "-p", dest="patterns", action="append_const",
                      const=PROCESS_RE, help="Color code by process (default)")
  parser.add_argument("--thread", "-t", dest="patterns", action="append_const",
                      const=THREAD_RE, help="Color code by thread")
  parser.add_argument("--regex", "-r", dest="patterns", action="append",
                      help="Color by matching regexp")
  parser.add_argument("input", nargs='?', action="store", default=None,
                      help="The file to colorize.  Defaults to stdin")
  args = parser.parse_args()

  if args.input:
    f = open(args.input, 'r')
  else:
    f = sys.stdin

  # If no patterns were specified, use the process pattern.
  if not args.patterns:
    args.patterns = [PROCESS_RE]

  # Define the identifier extractor.  It should be in group 1.
  patterns = []
  for pattern in args.patterns:
    regex = re.compile(pattern)
    if not regex.groups:
      # if there's no group, wrap the pattern
      regex = re.compile(r'^.*(' + pattern + r').*$')
    patterns.append(regex)

  assigner = ColorAssigner(BASE_COLORS);

  while True:
    line = f.readline()
    if not line:
      break

    line = line.strip()
    matched = False
    for line_re in patterns:
      m = line_re.match(line)
      if m:
        matched = True
        task_id = m.group(1)
        color = assigner.get_bg_color(task_id)

        # Use join to avoid python putting a space between each value being
        # printed.
        print ''.join([color, line, RESET_BG, RESET_FG])
        sys.stdout.flush()
        break

    if not matched:
      print line
      sys.stdout.flush()

if __name__ == '__main__':
  main()