from enum import Enum
from typing import Optional

from typing_extensions import NotRequired, TypedDict

from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams

DD_MAX_BATCH_SIZE = 1000


class DataDogStatus(str, Enum):
    INFO = "info"
    WARN = "warning"
    ERROR = "error"


DatadogPayload = TypedDict(
    "DatadogPayload",
    {
        "ddsource": str,
        "ddtags": str,
        "hostname": str,
        "message": str,
        "service": str,
        "status": str,
        "dd.trace_id": NotRequired[str],
        "dd.span_id": NotRequired[str],
    },
    total=False,
)


class DD_ERRORS(Enum):
    DATADOG_413_ERROR = "Datadog API Error - Payload too large (batch is above 5MB uncompressed). If you want this logged either disable request/response logging or set `DD_BATCH_SIZE=50`"


class DatadogInitParams(StandardCustomLoggerInitParams):
    """
    Params for initializing a DataDog logger on litellm
    """

    pass


class DatadogProxyFailureHookJsonMessage(TypedDict, total=False):
    exception: str
    error_class: str
    status_code: Optional[int]
    traceback: str
    user_api_key_dict: dict