#  Copyright (c) ZenML GmbH 2020. All Rights Reserved.
#
#  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:
#
#       https://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.
"""Logger implementation."""

import logging
import os
import re
import sys
from contextvars import ContextVar
from typing import Any, Dict, Optional

from rich.traceback import install as rich_tb_install

from zenml.constants import (
    ENABLE_RICH_TRACEBACK,
    ENV_ZENML_LOGGING_COLORS_DISABLED,
    ENV_ZENML_SUPPRESS_LOGS,
    ZENML_LOGGING_VERBOSITY,
    ZENML_STORAGE_LOGGING_VERBOSITY,
    handle_bool_env_var,
)
from zenml.enums import LoggingLevels

ZENML_LOGGING_COLORS_DISABLED = handle_bool_env_var(
    ENV_ZENML_LOGGING_COLORS_DISABLED, False
)

step_names_in_console: ContextVar[bool] = ContextVar(
    "step_names_in_console", default=False
)

_original_stdout_write: Optional[Any] = None
_original_stderr_write: Optional[Any] = None
_stdout_wrapped: bool = False
_stderr_wrapped: bool = False


class _ZenMLStdoutStream:
    """Stream that writes to the original stdout, bypassing the ZenML wrapper.

    This ensures console logging doesn't trigger the LoggingContext wrapper,
    preventing duplicate log entries in stored logs.
    """

    def write(self, text: str) -> Any:
        """Write text to the original stdout.

        Args:
            text: The text to write.

        Returns:
            The number of characters written.
        """
        if _original_stdout_write:
            return _original_stdout_write(text)
        return sys.stdout.write(text)

    def flush(self) -> None:
        """Flush the stdout buffer."""
        sys.stdout.flush()


def get_logger(logger_name: str) -> logging.Logger:
    """Main function to get logger name,.

    Args:
        logger_name: Name of logger to initialize.

    Returns:
        A logger object.
    """
    return logging.getLogger(logger_name)


def _add_step_name_to_message(message: str) -> str:
    """Adds the step name to the message.

    Args:
        message: The message to add the step name to.

    Returns:
        The message with the step name added.
    """
    try:
        if step_names_in_console.get():
            from zenml.steps import get_step_context

            step_context = get_step_context()

            if step_context and message not in ["\n", ""]:
                # For progress bar updates (with \r), inject the step name after the \r
                if "\r" in message:
                    message = message.replace(
                        "\r", f"\r[{step_context.step_name}] "
                    )
                else:
                    message = f"[{step_context.step_name}] {message}"
    except Exception:
        # If we can't get step context, just use the original message
        pass

    return message


class ConsoleFormatter(logging.Formatter):
    """Formats logs according to custom specifications."""

    grey: str = "\x1b[90m"
    white: str = "\x1b[37m"
    pink: str = "\x1b[35m"
    green: str = "\x1b[32m"
    yellow: str = "\x1b[33m"
    red: str = "\x1b[31m"
    cyan: str = "\x1b[1;36m"
    bold_red: str = "\x1b[31;1m"
    purple: str = "\x1b[38;5;105m"
    blue: str = "\x1b[34m"
    reset: str = "\x1b[0m"

    def _get_format_template(self, record: logging.LogRecord) -> str:
        """Get the format template based on the logging level.

        Args:
            record: The log record to format.

        Returns:
            The format template string.
        """
        if get_logging_level() == LoggingLevels.DEBUG:
            return "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
        else:
            return "%(message)s"

    COLORS: Dict[LoggingLevels, str] = {
        LoggingLevels.DEBUG: grey,
        LoggingLevels.INFO: white,
        LoggingLevels.WARN: yellow,
        LoggingLevels.ERROR: red,
        LoggingLevels.CRITICAL: bold_red,
    }

    def format(self, record: logging.LogRecord) -> str:
        """Converts a log record to a (colored) string.

        Args:
            record: LogRecord generated by the code.

        Returns:
            A string formatted according to specifications.
        """
        format_template = self._get_format_template(record)

        message = record.getMessage()
        try:
            if step_names_in_console.get():
                message = _add_step_name_to_message(message)
        except Exception:
            pass

        modified_record = logging.LogRecord(
            name=record.name,
            level=record.levelno,
            pathname=record.pathname,
            lineno=record.lineno,
            msg=message,
            args=(),
            exc_info=record.exc_info,
        )

        if ZENML_LOGGING_COLORS_DISABLED:
            formatter = logging.Formatter(format_template)
            return formatter.format(modified_record)
        else:
            log_fmt = (
                self.COLORS[LoggingLevels(record.levelno)]
                + format_template
                + self.reset
            )
            formatter = logging.Formatter(log_fmt)
            formatted_message = formatter.format(modified_record)
            quoted_groups = re.findall("`([^`]*)`", formatted_message)
            for quoted in quoted_groups:
                formatted_message = formatted_message.replace(
                    "`" + quoted + "`",
                    self.reset
                    + self.purple
                    + quoted
                    + self.COLORS.get(LoggingLevels(record.levelno)),
                )

            url_pattern = r"http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+"
            urls = re.findall(url_pattern, formatted_message)
            for url in urls:
                formatted_message = formatted_message.replace(
                    url,
                    self.reset
                    + self.blue
                    + url
                    + self.COLORS.get(LoggingLevels(record.levelno)),
                )
            return formatted_message


def get_logging_level() -> LoggingLevels:
    """Get logging level from the env variable.

    Returns:
        The logging level.

    Raises:
        KeyError: If the logging level is not found.
    """
    verbosity = ZENML_LOGGING_VERBOSITY.upper()
    if verbosity not in LoggingLevels.__members__:
        raise KeyError(
            f"Verbosity must be one of {list(LoggingLevels.__members__.keys())}"
        )

    if ZENML_STORAGE_LOGGING_VERBOSITY is not None:
        get_logger(__name__).warning(
            "The ZENML_STORAGE_LOGGING_VERBOSITY is no longer supported. "
            "Please use the ZENML_LOGGING_VERBOSITY instead."
        )

    return LoggingLevels[verbosity]


def set_root_verbosity() -> None:
    """Set the root verbosity."""
    level = get_logging_level()
    if level != LoggingLevels.NOTSET:
        if ENABLE_RICH_TRACEBACK:
            rich_tb_install(show_locals=(level == LoggingLevels.DEBUG))

        logging.root.setLevel(level=level.value)
        get_logger(__name__).debug(
            f"Logging set to level: {logging.getLevelName(level.value)}"
        )
    else:
        logging.disable(sys.maxsize)
        logging.getLogger().disabled = True
        get_logger(__name__).debug("Logging NOTSET")


def _wrapped_write(original_write: Any, stream_name: str) -> Any:
    """Wrap stdout/stderr write method to route logs to LoggingContext.

    Args:
        original_write: The original write method.
        stream_name: The name of the stream.

    Returns:
        The wrapped write method.
    """

    def wrapped_write(text: str) -> Any:
        """Write method that routes logs through LoggingContext.

        Args:
            text: The text to write.

        Returns:
            The result of the original write method.
        """
        from zenml.utils.logging_utils import LoggingContext

        level = logging.INFO if stream_name == "stdout" else logging.ERROR

        if logging.root.isEnabledFor(level):
            record = logging.LogRecord(
                name=stream_name,
                level=level,
                pathname="",
                lineno=0,
                msg=text,
                args=(),
                exc_info=None,
                func="",
            )
            LoggingContext.emit(record)

        return original_write(text)

    return wrapped_write


def wrap_stdout_stderr() -> None:
    """Wrap stdout and stderr write methods to route through LoggingContext."""
    global _stdout_wrapped, _stderr_wrapped
    global _original_stdout_write, _original_stderr_write

    if not _stdout_wrapped:
        _original_stdout_write = getattr(sys.stdout, "write")
        setattr(
            sys.stdout,
            "write",
            _wrapped_write(_original_stdout_write, "stdout"),
        )
        _stdout_wrapped = True

    if not _stderr_wrapped:
        _original_stderr_write = getattr(sys.stderr, "write")
        setattr(
            sys.stderr,
            "write",
            _wrapped_write(_original_stderr_write, "stderr"),
        )
        _stderr_wrapped = True


class ZenMLLoggingHandler(logging.Handler):
    """Custom handler that routes logs through LoggingContext."""

    def emit(self, record: logging.LogRecord) -> None:
        """Emit a log record through LoggingContext.

        Args:
            record: The log record to emit.
        """
        from zenml.utils.logging_utils import LoggingContext

        LoggingContext.emit(record)


def get_console_handler() -> logging.Handler:
    """Get console handler that writes to original stdout.

    Returns:
        A console handler.
    """
    handler = logging.StreamHandler(_ZenMLStdoutStream())
    handler.setFormatter(ConsoleFormatter())
    return handler


def get_zenml_handler() -> logging.Handler:
    """Get ZenML handler that routes logs through LoggingContext.

    Returns:
        A ZenML handler.
    """
    return ZenMLLoggingHandler()


def init_logging() -> None:
    """Initialize the logging system."""
    set_root_verbosity()
    wrap_stdout_stderr()

    # Add both handlers to the root logger
    root_logger = logging.getLogger()

    # Console handler - writes to original stdout
    root_logger.addHandler(get_console_handler())

    # ZenML handler - routes through LoggingContext
    root_logger.addHandler(get_zenml_handler())

    # Mute tensorflow cuda warnings
    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

    # logging capture warnings
    logging.captureWarnings(True)

    # Enable logs if environment variable SUPPRESS_ZENML_LOGS is not set to True
    suppress_zenml_logs: bool = handle_bool_env_var(
        ENV_ZENML_SUPPRESS_LOGS, True
    )
    if suppress_zenml_logs:
        # suppress logger info messages
        suppressed_logger_names = [
            "urllib3",
            "azure.core.pipeline.policies.http_logging_policy",
            "grpc",
            "requests",
            "kfp",
            "tensorflow",
        ]
        for logger_name in suppressed_logger_names:
            logging.getLogger(logger_name).setLevel(logging.WARNING)

        # disable logger messages
        disabled_logger_names = [
            "rdbms_metadata_access_object",
            "backoff",
            "segment",
        ]
        for logger_name in disabled_logger_names:
            logging.getLogger(logger_name).setLevel(logging.WARNING)
            logging.getLogger(logger_name).disabled = True
