#!/usr/bin/env python3
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import argparse
import collections
from collections.abc import Callable, Iterator, Sequence
import functools
import io
import sys
import time
from typing import cast

from oslo_serialization import jsonutils
from oslo_utils import importutils

from oslo_log.formatters import JSONLogRecord
from oslo_log import log

termcolor = importutils.try_import('termcolor')


_USE_COLOR = False
DEFAULT_LEVEL_KEY = 'levelname'
DEFAULT_TRACEBACK_KEY = 'traceback'


def main() -> None:
    global _USE_COLOR
    args = parse_args()
    _USE_COLOR = args.color
    formatter = functools.partial(
        console_format,
        args.prefix,
        args.locator,
        loggers=args.loggers,
        levels=args.levels,
        level_key=args.levelkey,
        traceback_key=args.tbkey,
    )
    if args.lines:
        # Read backward until we find all of our newline characters
        # or reach the beginning of the file
        args.file.seek(0, 2)
        newlines = 0
        pos = args.file.tell()
        while newlines <= args.lines and pos > 0:
            pos = pos - 1
            args.file.seek(pos)
            if args.file.read(1) == '\n':
                newlines = newlines + 1
    try:
        for line in reformat_json(args.file, formatter, args.follow):
            print(line)
    except KeyboardInterrupt:
        sys.exit(0)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "file",
        nargs='?',
        default=sys.stdin,
        type=argparse.FileType(),
        help="JSON log file to read from (if not provided"
        " standard input is used instead)",
    )
    parser.add_argument(
        "--prefix",
        default='%(asctime)s.%(msecs)03d %(process)s %(levelname)s %(name)s',
        help="Message prefixes",
    )
    parser.add_argument(
        "--locator",
        default='[%(funcname)s %(pathname)s:%(lineno)s]',
        help="Locator to append to DEBUG records",
    )
    parser.add_argument(
        "--levelkey",
        default=DEFAULT_LEVEL_KEY,
        help="Key in the JSON record where the level is held",
    )
    parser.add_argument(
        "--tbkey",
        default=DEFAULT_TRACEBACK_KEY,
        help="Key in the JSON record where the traceback/exception is held",
    )
    parser.add_argument(
        "-c",
        "--color",
        action='store_true',
        default=False,
        help="Color log levels (requires `termcolor`)",
    )
    parser.add_argument(
        "-f",
        "--follow",
        action='store_true',
        default=False,
        help="Continue parsing new data until KeyboardInterrupt",
    )
    parser.add_argument(
        "-n",
        "--lines",
        required=False,
        type=int,
        help="Last N number of records to view."
        " (May show less than N records when used"
        " in conjuction with --loggers or --levels)",
    )
    parser.add_argument(
        "--loggers",
        nargs='*',
        default=[],
        help="only return results matching given logger(s)",
    )
    parser.add_argument(
        "--levels",
        nargs='*',
        default=[],
        help="Only return lines matching given log level(s)",
    )
    args = parser.parse_args()
    if args.color and not termcolor:
        raise ImportError(
            "Coloring requested but `termcolor` is not importable"
        )
    return args


def colorise(key: str, text: str | None = None) -> str:
    if text is None:
        text = key
    if not _USE_COLOR:
        return text
    colors = {
        'exc': ('red', ['reverse', 'bold']),
        'FATAL': ('red', ['reverse', 'bold']),
        'ERROR': ('red', ['bold']),
        'WARNING': ('yellow', ['bold']),
        'WARN': ('yellow', ['bold']),
        'INFO': ('white', ['bold']),
    }
    color, attrs = colors.get(key, ('', []))
    if color:
        return cast(str, termcolor.colored(text, color=color, attrs=attrs))
    return text


def warn(prefix: str, msg: object) -> str:
    return "{}: {}".format(colorise('exc', prefix), msg)


def reformat_json(
    fh: io.StringIO,
    formatter: Callable[..., Iterator[str]],
    follow: bool = False,
) -> Iterator[str]:
    # using readline allows interactive stdin to respond to every line
    while True:
        line = fh.readline()
        if not line:
            if follow:
                time.sleep(0.1)
                continue
            else:
                break
        line = line.strip()
        if not line:
            continue
        try:
            record = jsonutils.loads(line)
        except ValueError:
            yield warn("Not JSON", line)
            continue
        yield from formatter(record)


def console_format(
    prefix: str,
    locator: str,
    record: JSONLogRecord,
    loggers: Sequence[str] = [],
    levels: Sequence[str] = [],
    level_key: str = DEFAULT_LEVEL_KEY,
    traceback_key: str = DEFAULT_TRACEBACK_KEY,
) -> Iterator[str]:
    # Provide an empty string to format-specifiers the record is
    # missing, instead of failing. Doesn't work for non-string
    # specifiers.
    record = collections.defaultdict(str, record)  # type: ignore
    # skip if the record doesn't match a logger we are looking at
    if loggers:
        name = record['name']
        if not any(name.startswith(n) for n in loggers):
            return

    if levels:
        if record.get(level_key) not in levels:
            return

    levelname = record.get(level_key)
    if levelname:
        record[level_key] = colorise(levelname)  # type: ignore

    try:
        prefix = prefix % record
    except TypeError:
        # Thrown when a non-string format-specifier can't be filled in.
        # Dict comprehension cleans up the output
        yield warn(
            'Missing non-string placeholder in record',
            {
                str(k): str(v) if isinstance(v, str) else v
                for k, v in record.items()
            },
        )
        return

    locator = ''
    if record.get('levelno', 100) <= log.DEBUG or levelname == 'DEBUG':
        locator = locator % record

    yield ' '.join(x for x in [prefix, record['message'], locator] if x)

    tb = record.get(traceback_key)
    if tb:
        if type(tb) is str:
            tb = tb.rstrip().split("\n")
        for tb_line in tb:  # type: ignore
            yield ' '.join([prefix, tb_line])


if __name__ == '__main__':
    main()
