# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

import importlib
import os
import sys
from datetime import datetime
from types import ModuleType
from typing import Any, Optional

from cosmos_framework.utils.log import LEVEL, get_datetime_format, get_machine_format, logger, make_new_logger

__enable_telemetry: bool = os.getenv("ENABLE_TELEMETRY", "FALSE").upper() == "TRUE"
__telemetry_backends: set[str] = set(os.getenv("TELEMETRY_BACKENDS", "logger").split(","))
__training_telemetry_module: Optional[ModuleType] = None
__provider = None


def import_training_telemetry() -> Optional[ModuleType]:
    """Lazy import of the training_telemetry package to avoid a hard dependency."""
    global __training_telemetry_module
    if __training_telemetry_module is not None:
        return __training_telemetry_module

    try:
        __training_telemetry_module = importlib.import_module("training_telemetry")
        return __training_telemetry_module
    except ImportError as e:
        logger.error(f"Telemetry is enabled but the `training_telemetry` package is not installed: {e}")
        return None


def set_telemetry_provider(local_path: str) -> Optional[Any]:
    """
    Set the telemetry provider if telemetry is enabled, and if we can import the necessary modules, otherwise return None.
    """
    global __enable_telemetry
    if not __enable_telemetry:
        logger.info("Training telemetry is disabled. Set ENABLE_TELEMETRY=true to enable it.")
        return None

    global __provider
    if __provider is not None:
        return __provider

    training_telemetry = import_training_telemetry()
    if training_telemetry is None:
        logger.error(
            "Training telemetry is enabled but the `training_telemetry` package is not installed. "
            "Set ENABLE_TELEMETRY=false to disable, or install the package."
        )
        __enable_telemetry = False
        return None

    rank = training_telemetry.get_rank()
    log_only_errors = rank != 0
    logger.info(
        f"Setting up telemetry provider, rank: {rank}, log only errors: {log_only_errors}, backends: {__telemetry_backends}"
    )
    backends = []
    for backend_name in __telemetry_backends:
        backend_name = backend_name.strip().lower()
        if not backend_name:
            continue
        logger.debug(f"Setting up telemetry backend: {backend_name}")
        if backend_name == "logger":
            backends.append(
                training_telemetry.LoggerBackendConfig(rank_aware=True, errors_only=log_only_errors),
            )
        elif backend_name == "file":
            backends.append(training_telemetry.FileBackendConfig(output_file_path=local_path + "/telemetry/events"))
        elif backend_name == "nvtx":
            backends.append(training_telemetry.NVTXBackendConfig())
        else:
            logger.error(f"Unknown telemetry backend will be ignored: {backend_name}")
    if not backends:
        logger.error("No telemetry backends configured, telemetry will be disabled")
        __enable_telemetry = False
        return None

    config = training_telemetry.TelemetryConfig(backends=backends)

    # Use a telemetry-specific logger because we don't want to report information that is not useful for telemetry
    # such as the file path to the library and the function name of the logging backend
    telemetry_logger = make_new_logger(depth=0)
    message_format = "<level>{level}</level>|training-telemetry] {message}"
    telemetry_logger.remove()
    telemetry_logger.add(
        sys.stdout,
        level=LEVEL,
        format=f"{get_datetime_format()}{get_machine_format()}{message_format}",
    )
    telemetry_logger.add(
        f"{local_path}/telemetry/stdout.log",
        encoding="utf8",
        level=LEVEL,
        format=f"{get_datetime_format()}{get_machine_format()}{message_format}",
        rotation="100 MB",
        enqueue=True,
    )
    __provider = training_telemetry.Provider.set_provider(config=config, logger=telemetry_logger)
    return __provider


def get_telemetry_recorder() -> Optional[Any]:
    """
    Get the telemetry recorder if it is available, otherwise return None.
    """
    global __provider
    return __provider.recorder if __provider else None


def get_timezone_name() -> str:
    """Get the timezone name, tzlocal is installed with the training-telemetry package, if it is not installed, use the TZ environment variable, otherwise use the local timezone"""
    try:
        # pyrefly: ignore  # import-error
        from tzlocal import get_localzone

        return str(get_localzone())
    except ImportError:
        tz_name = os.environ.get("TZ")
        if tz_name is None:
            tz_name = datetime.now().astimezone().tzname()
        return tz_name


def get_checkpoint_strategy(checkpoint_config: dict[str, Any]) -> str:
    """
    Return the checkpoint strategy, sync or async, based on the checkpoint config.

    FIXME: it seems that the config only reports the type of checkpoint, sync or async,
    for the cosmos distributed checkpointer, but I don't want a dependency  on it here
    so I'm looking at the class name. It seems that the default Checkpointer uses an async thread
    all the time, so the default is async. The FIX should involve either having a config that is
    always followed by all checkpointers, or adding a function to the base checkpointer class,
    indicating if checkpoint is sync or async, but this means that the checkpoint class could
    only be retrieved after the checkpointer is created, and not from the config.
    """
    if checkpoint_config.type is not None and "DistributedCheckpointer" in checkpoint_config.type["_target_"].__name__:
        return "async" if checkpoint_config.dcp_async_mode_enabled else "sync"
    else:
        return "async"