# SPDX-License-Identifier: Apache-2.0
"""Configuration for OneLogger."""

from enum import Enum
from pathlib import Path
from typing import Callable, Dict, Literal, Optional, Union

from pydantic import BaseModel, ConfigDict, model_validator
from typing_extensions import Self

from nv_one_logger.api.telemetry_config import TelemetryConfig
from nv_one_logger.core.attributes import AttributeValue
from nv_one_logger.core.internal.utils import evaluate_value


class OneLoggerErrorHandlingStrategy(Enum):
    """Enum for the error handling strategy for OneLogger.

    This enum determines what happens when OneLogger encounters a fatal error (e.g., an exception in the instrumentation code or
    a problem with the OneLogger state). This does NOT affect handling of errors occuring when communicating with the telemetry backends
    (i.e., exporter failures, which are handled by the Recorder). Rather, this is about handling user errors when configuring OneLogger
    or bugs in telemetry code (e.g., assertion/invariant violations or hitting an inconsistent state).

    Our recommendation is to use PROPAGATE_EXCEPTIONS as it ensures that you get maximum visibility into the errors in the instrumentation code.
    But if you prefer to treat telemetry as a non-critical part of your application, you can use DISABLE_QUIETLY_AND_REPORT_METRIC_ERROR.

    Read the docstrings for DefaultRecorder for more details on how errors from exporters (e.g., communication errors with telemetry
    backends) are handled when using that recorder.
    """

    """Propagate the exceptions to the caller.

    Use this strategy if you are OK with instrumentation exceptions to crash the application (or are willing to catch and handle such exceptions).
    Note that all of exceptions from One Logger will be instances of `OneLoggerError` or a subclass of it. This allows you to identify those exceptions
    and react to them accordingly.
    The advantage of this strategy is that you get maximum visibility into the errors in the instrumentation code.
    """
    PROPAGATE_EXCEPTIONS = "propagate_exceptions"

    """Disable OneLogger silently and report metric errors to the telemetry backends.

    With this option, if instrumentation code encounters any errors, the library catches/suppresses the exception letting the
    application continue running, logs a single error to the application logs, and informs the telemtry backends that
    the telemetry data has errors.
    """
    DISABLE_QUIETLY_AND_REPORT_METRIC_ERROR = "disable_quietly_and_report_metric_error"


class LoggerConfig(BaseModel):
    """Configuration for how OneLogger logs its messages and errors."""

    log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR"] = "INFO"

    log_format: str = "%(name)s - %(levelname)s - %(message)s"

    # Path to the file where OneLogger INFO logs its messages.
    log_file_path_for_info: Union[Path, str] = "onelogger.log"

    # Path to the file where OneLogger ERROR logs its messages.
    log_file_path_for_err: Union[Path, str] = "onelogger.err"

    @model_validator(mode="after")
    def validate_config(self) -> Self:
        """Validate the logger config."""
        if self.log_level not in {"DEBUG", "INFO", "WARNING", "ERROR"}:
            raise ValueError(f"log_level must be one of {'DEBUG', 'INFO', 'WARNING', 'ERROR'}, got {self.log_level}")
        return self


class OneLoggerConfig(BaseModel):
    """Configuration for OneLogger."""

    # Pydantic model configuration to allow arbitrary types in field definitions.
    # This is necessary because the telemetry_config field uses a Protocol type (TelemetryConfig),
    # which is a custom type that Pydantic cannot automatically generate a schema for.
    # By setting arbitrary_types_allowed=True,
    # we tell Pydantic to accept any type for this field and skip schema validation for it.
    #
    model_config = ConfigDict(arbitrary_types_allowed=True)

    # The unique name for application. This name is used to identify the telemetry data related to various executions of
    # the same application in the OneLogger system (over time, across different machines/clusters, and across
    # different versions of the application).
    application_name: str

    # Number of processes participating in the telemetry collection.
    # This is a fundamental configuration that affects the entire application, not just training telemetry.
    world_size_or_fn: Union[int, Callable[[], int]]

    @property
    def world_size(self) -> int:
        """Number of processes participating in the telemetry collection."""
        return evaluate_value(self.world_size_or_fn)

    # session_tag (or callable to generate the tag). session_tag is used to used to identify jobs that together contribute
    # to the same task. This means the jobs are "logically" part of a single larger job (e.g.,
    # a long running job that is split into multiple jobs due to resource constraints or resuming after a failure).
    session_tag_or_fn: Union[str, Callable[[], str]]

    @property
    def session_tag(self) -> str:
        """Get the session tag.

        Returns:
            str: The evaluated session tag value.
        """
        return evaluate_value(self.session_tag_or_fn)

    # Flag (or callable to return flag) that indicates if this is a baseline run for comparison purposes.
    # A baseline run is a run that is used to set a performance baseline for future runs.
    is_baseline_run_or_fn: Union[bool, Callable[[], bool]] = False

    @property
    def is_baseline_run(self) -> bool:
        """Get the baseline run flag.

        Returns:
            bool: The evaluated baseline run flag value.
        """
        return evaluate_value(self.is_baseline_run_or_fn)

    # Custom metadata to be logged with the training telemetry data.
    # This metadata will be logged as-is, without any modification as an
    # attribute of the APPLICATION span.
    custom_metadata: Optional[Dict[str, AttributeValue]] = None

    # The strategy to use for handling errors in the instrumentation code.
    # By default, we will propagate the exceptions to the caller. This means that if there is an error in the instrumentation
    # code, it will crash the application.
    # You can choose a different strategy if you want to treat telemetry as a non-critical part of your application.
    # See the enum docstring for more details on each strategy and for our recommendations on how to set this value.
    error_handling_strategy: OneLoggerErrorHandlingStrategy = OneLoggerErrorHandlingStrategy.PROPAGATE_EXCEPTIONS

    # Whether to enable logging for the current rank in distributed training.
    # This controls whether OneLogger is enabled for the current process/rank.
    # In distributed training scenarios, you can set this to False for ranks where you don't want logging.
    enable_for_current_rank: bool = True

    # Configuration for the logger used for logging messages and errors from the telemetry code.
    logger_config: LoggerConfig = LoggerConfig()

    # Version (or callable to return version) of the data schema used for summarizing metrics.
    # If the schema of the data you collect changes over time, you can use this value to
    # keep track of which schema version is used for which run.
    summary_data_schema_version: Union[str, str] = "1.0.0"

    # Telemetry-specific configuration.
    # This field contains telemetry-specific settings and parameters.
    # It is optional and can be None if no telemetry is needed.
    #
    # Note: This field uses the TelemetryConfig Protocol type, which is why we need
    # arbitrary_types_allowed=True in the model_config above. The actual value must be
    # a concrete class that implements the TelemetryConfig protocol interface.
    telemetry_config: Optional[TelemetryConfig] = None

    @model_validator(mode="after")
    def validate_config(self) -> Self:
        """Validate the OneLogger configuration.

        This validator ensures that:
        - application_name is not empty
        - world_size is set to a positive value
        - custom_metadata keys are valid strings (if provided)
        - telemetry_config implements the TelemetryConfig protocol (if provided)

        Returns:
            OneLoggerConfig: The validated configuration.

        Raises:
            ValueError: If the configuration is invalid.
        """
        if not self.application_name or not self.application_name.strip():
            raise ValueError("application_name cannot be empty or whitespace-only")

        if self.world_size <= 0:
            raise ValueError("world_size must be set to a positive value")

        if self.custom_metadata is not None:
            for key in self.custom_metadata.keys():
                if not key.strip():
                    raise ValueError("custom_metadata keys must be non-empty strings")

        return self
