# Copyright (C) 2024 Bloomberg LP
#
# 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 logging
from enum import Enum
from types import TracebackType
from typing import Any, Dict, Optional, Tuple, Type, Union

from google.protobuf import text_format
from google.protobuf.message import Message

from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Digest

Exc = Union[
    bool,
    Tuple[Type[BaseException], BaseException, Optional[TracebackType]],
    Tuple[None, None, None],
    BaseException,
]

Tags = Dict[str, Any]


def _str_escape(s: str) -> str:
    return str(s).replace('"', r"\"")


def _format_log_tag_value(value: Any) -> Any:
    if value is None:
        return '""'
    elif isinstance(value, int):
        return value
    elif isinstance(value, float):
        return f"{value:.2f}"
    elif isinstance(value, Digest):
        return f'"{value.hash}/{value.size_bytes}"'
    elif isinstance(value, Message):
        return f'"{_str_escape(text_format.MessageToString(value, as_one_line=True))}"'
    elif isinstance(value, Enum):
        return value.name
    else:
        return f'"{_str_escape(value)}"'


def _format_log_tags(tags: Optional[Tags]) -> str:
    if not tags:
        return ""
    return "".join([f" {key}={_format_log_tag_value(value)}" for key, value in sorted(tags.items())])


class BuildgridLogger:
    def __init__(self, logger: logging.Logger) -> None:
        """
        The buildgrid logger is a helper utility wrapped around a standard logger instance.
        It allows placing key=value strings at the end of log lines, reducing boilerplate in
        displaying values and adding standardization to our log lines. Within each logging method,
        tags may be added by setting the "tags" argument.

        Each logger is set to log at stacklevel=2 such that function names and source line numbers
        show the line at which this utility is invoked.

        Special encoding rules for tag values:
         - int: reported as is. `value=1`
         - float: rounded to the nearest two decimals. `value=1.23`
         - Digest: unpacked as hash/size. `value=deadbeef/123`
         - proto.Message: text_format, escaped, and quoted. `value="blob_digests { hash: \"deadbeef\" }"`
         - Enum: attribute name is used. `value=OK`
         - others: converted to str, escaped, and quoted. `value="foo: \"bar\""`

        Encoding is only performed if logging is enabled for that level.
        """
        self._logger = logger

    def is_enabled_for(self, level: int) -> bool:
        return self._logger.isEnabledFor(level)

    def debug(self, msg: Any, *, exc_info: Optional[Exc] = None, tags: Optional[Tags] = None) -> None:
        if self._logger.isEnabledFor(logging.DEBUG):
            self._logger.debug(str(msg) + _format_log_tags(tags), exc_info=exc_info, stacklevel=2)

    def info(self, msg: Any, *, exc_info: Optional[Exc] = None, tags: Optional[Tags] = None) -> None:
        if self._logger.isEnabledFor(logging.INFO):
            self._logger.info(str(msg) + _format_log_tags(tags), exc_info=exc_info, stacklevel=2)

    def warning(self, msg: Any, *, exc_info: Optional[Exc] = None, tags: Optional[Tags] = None) -> None:
        if self._logger.isEnabledFor(logging.WARNING):
            self._logger.warning(str(msg) + _format_log_tags(tags), exc_info=exc_info, stacklevel=2)

    def error(self, msg: Any, *, exc_info: Optional[Exc] = None, tags: Optional[Tags] = None) -> None:
        if self._logger.isEnabledFor(logging.ERROR):
            self._logger.error(str(msg) + _format_log_tags(tags), exc_info=exc_info, stacklevel=2)

    def exception(self, msg: Any, *, exc_info: Optional[Exc] = True, tags: Optional[Tags] = None) -> None:
        if self._logger.isEnabledFor(logging.ERROR):
            # Note we call error here instead of exception.
            # logger.exception is a helper around calling error with exc_info defaulting to True.
            # On python<3.11 that helper causes the stacklevel to report incorrectly.
            self._logger.error(str(msg) + _format_log_tags(tags), exc_info=exc_info, stacklevel=2)


def buildgrid_logger(name: str) -> BuildgridLogger:
    return BuildgridLogger(logging.getLogger(name))
