# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 json
import logging
import os
import tempfile

from dynamo._core import log_message


class LogHandler(logging.Handler):
    """
    Custom logging handler that sends log messages to the Rust env_logger
    """

    def emit(self, record):
        """
        Emit a log record
        """
        log_entry = self.format(record)
        if record.funcName == "<module>":
            module_path = record.module
        else:
            module_path = f"{record.module}.{record.funcName}"
        log_message(
            record.levelname.lower(),
            log_entry,
            module_path,
            record.pathname,
            record.lineno,
        )


# Configure the Python logger to use the NimLogHandler
def configure_logger(service_name: str | None, worker_id: int | None):
    """
    Called once to configure the Python logger to use the LogHandler
    """
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    handler = LogHandler()

    # Simple formatter without date and level info since it's already provided by Rust
    formatter = logging.Formatter("%(message)s")
    formatter_prefix = construct_formatter_prefix(service_name, worker_id)
    if len(formatter_prefix) != 0:
        formatter = logging.Formatter(f"[{formatter_prefix}] %(message)s")

    handler.setFormatter(formatter)
    logger.addHandler(handler)


def construct_formatter_prefix(service_name: str | None, worker_id: int | None) -> str:
    tmp = ""
    if service_name is not None:
        tmp += f"{service_name}"

    if worker_id is not None:
        tmp += f":{worker_id}"

    return tmp.strip()


def configure_dynamo_logging(
    service_name: str | None = None, worker_id: int | None = None
):
    """
    A single place to configure logging for Dynamo.
    """
    # First, remove any existing handlers to avoid duplication
    root_logger = logging.getLogger()
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)

    # Configure the logger with Dynamo's handler
    configure_logger(service_name, worker_id)

    # map the DYN_LOG variable to a logging level
    dyn_var = os.environ.get("DYN_LOG", "info")
    dyn_level = log_level_mapping(dyn_var)

    # configure inference engine loggers
    if not get_bool_env_var("DYN_SKIP_VLLM_LOG_FORMATTING"):
        configure_vllm_logging(dyn_level)
    if not get_bool_env_var("DYN_SKIP_SGLANG_LOG_FORMATTING"):
        configure_sglang_logging(dyn_level)
    if not get_bool_env_var("DYN_SKIP_TRTLLM_LOG_FORMATTING"):
        configure_trtllm_logging(dyn_level)

    # loggers that should be configured to ERROR
    error_loggers = ["tag"]
    for logger_name in error_loggers:
        logger = logging.getLogger(logger_name)
        logger.handlers = []
        logger.setLevel(logging.ERROR)
        logger.propagate = True


def log_level_mapping(level: str) -> int:
    """
    The DYN_LOG variable is set using "debug" or "trace" or "info.
    This function maps those to the appropriate logging level and defaults to INFO
    if the variable is not set or a bad value.
    """
    if level == "debug":
        return logging.DEBUG
    elif level == "info":
        return logging.INFO
    elif level == "warn" or level == "warning":
        return logging.WARNING
    elif level == "error":
        return logging.ERROR
    elif level == "critical":
        return logging.CRITICAL
    elif level == "trace":
        return logging.INFO
    else:
        return logging.INFO


def configure_sglang_logging(dyn_level: int):
    """
    SGLang allows us to create a custom logging config file
    """

    sglang_level = logging.getLevelName(dyn_level)

    sglang_config = {
        "formatters": {"simple": {"format": "%(message)s"}},
        "handlers": {
            "dynamo": {
                "class": "dynamo.runtime.logging.LogHandler",
                "formatter": "simple",
                "level": sglang_level,
            }
        },
        "loggers": {
            "sglang": {
                "handlers": ["dynamo"],
                "level": sglang_level,
                "propagate": False,
            }
        },
        "version": 1,
        "disable_existing_loggers": False,
    }

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
        json.dump(sglang_config, f)
        os.environ["SGLANG_LOGGING_CONFIG_PATH"] = f.name


def configure_vllm_logging(dyn_level: int):
    """
    vLLM requires a logging config file to be set in the environment.
    This function creates a temporary file with the VLLM logging config and sets the
    VLLM_LOGGING_CONFIG_PATH environment variable to the path of the file.
    """

    os.environ["VLLM_CONFIGURE_LOGGING"] = "1"
    vllm_level = logging.getLevelName(dyn_level)

    # Create a temporary config file for VLLM
    vllm_config = {
        "formatters": {"simple": {"format": "%(message)s"}},
        "handlers": {
            "dynamo": {
                "class": "dynamo.runtime.logging.LogHandler",
                "formatter": "simple",
                "level": vllm_level,
            }
        },
        "loggers": {
            "vllm": {"handlers": ["dynamo"], "level": vllm_level, "propagate": False}
        },
        "version": 1,
        "disable_existing_loggers": False,
    }

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
        json.dump(vllm_config, f)
        os.environ["VLLM_LOGGING_CONFIG_PATH"] = f.name


def map_dyn_log_to_tllm_level(dyn_log_value: str) -> str:
    """
    Map DYN_LOG string value to TensorRT-LLM log level.

    Args:
        dyn_log_value: The DYN_LOG environment variable value (e.g., "debug", "info,module::path=trace")

    Returns:
        The corresponding TLLM_LOG_LEVEL value (e.g., "VERBOSE", "INFO")
    """
    # Extract the base level (handle cases like "debug,module::path=trace")
    base_level = dyn_log_value.lower().split(",")[0].strip()

    # Map DYN_LOG levels to TLLM_LOG_LEVEL
    level_mapping = {
        "debug": "DEBUG",
        "info": "INFO",
        "warn": "WARNING",
        "warning": "WARNING",
        "error": "ERROR",
        "critical": "ERROR",
        "trace": "TRACE",
    }

    return level_mapping.get(base_level, "INFO")


def configure_trtllm_logging(dyn_level: int):
    # Only set TLLM_LOG_LEVEL if it's not already set
    # This allows users to override it explicitly if needed
    if "TLLM_LOG_LEVEL" not in os.environ:
        dyn_level_name = logging.getLevelName(dyn_level)
        tllm_level = map_dyn_log_to_tllm_level(dyn_level_name)
        os.environ["TLLM_LOG_LEVEL"] = tllm_level
        logging.debug(f"Set TLLM_LOG_LEVEL to {tllm_level} based on DYN_LOG")


def get_bool_env_var(name: str, default: str = "false") -> bool:
    value = os.getenv(name, default)
    value = value.lower()

    truthy_values = ("true", "1")
    falsy_values = ("false", "0")

    if (value not in truthy_values) and (value not in falsy_values):
        logging.warning(
            f"The environment variable {name} has an unrecognized value={value}, treating as false"
        )

    return value in truthy_values
