import os
import re
import secrets
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from datetime import datetime as dt
from typing import Final, Literal, Protocol, cast, runtime_checkable

from pydantic import BaseModel

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
    LITELLM_PROXY_MASTER_KEY_ALIAS,
    LITELLM_TRUNCATED_PAYLOAD_FIELD,
    LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
    REDACTED_BY_LITELM_STRING,
)
from litellm.constants import (
    MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB,
)
from litellm.litellm_core_utils.core_helpers import (
    get_litellm_metadata_from_kwargs,
    reconstruct_model_name,
)
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call
from litellm.litellm_core_utils.litellm_logging import (
    coerce_model_access_groups,
    is_valid_sha256_hash,
    request_model_access_groups_from_litellm_params,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsRouterMetadata
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.proxy.utils import PrismaClient, hash_token
from litellm.types.utils import (
    CallTypes,
    CostBreakdown,
    StandardLoggingGuardrailInformation,
    StandardLoggingMCPToolCall,
    StandardLoggingModelInformation,
    StandardLoggingPayload,
    StandardLoggingPayloadErrorInformation,
    StandardLoggingVectorStoreRequest,
    VectorStoreSearchResponse,
)
from litellm.utils import get_end_user_id_for_cost_tracking


def _get_max_string_length_prompt_in_db() -> int:
    """
    Resolve prompt truncation threshold at runtime so values loaded later via
    proxy config environment_variables are honored.
    """
    max_length_str: Final = os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB")
    if max_length_str is None:
        return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB
    try:
        return int(max_length_str)
    except (TypeError, ValueError):
        return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB


def _is_master_key(api_key: str | None, _master_key: str | None) -> bool:
    """
    Raw-only constant-time master-key comparison. The hashed form is never
    considered equivalent — only the raw master-key string matches.
    """
    if _master_key is None or api_key is None:
        return False
    return secrets.compare_digest(api_key, _master_key)


_HASHED_JWT_RE = re.compile(r"hashed-jwt-[a-fA-F0-9]{64}")


def _is_non_secret_key_value(value: str) -> bool:
    return (
        value == LITELLM_PROXY_MASTER_KEY_ALIAS
        or is_valid_sha256_hash(value)
        or _HASHED_JWT_RE.fullmatch(value) is not None
    )


def _redact_logged_api_key(value: str | None, *, already_redacted: bool = False) -> str | None:
    if not isinstance(value, str) or not value:
        return None
    stripped: Final = re.sub(r"(?i)^bearer ", "", value)
    if not stripped:
        return None
    if already_redacted and _is_non_secret_key_value(stripped):
        return stripped
    return hash_token(stripped)


def _get_router_metadata_for_spend_log(
    metadata: Mapping[str, object] | None,
    requested_model: str | None,
    selected_model: str | None,
    selected_provider: str | None,
    router_correlation_id: str | None,
) -> SpendLogsRouterMetadata | None:
    model_info: Final = metadata.get("model_info") if metadata is not None else None
    if not isinstance(model_info, Mapping) or model_info.get("internal_router_model") is not True:
        return None
    return SpendLogsRouterMetadata(
        requested_model=requested_model or None,
        selected_model=selected_model or None,
        selected_provider=selected_provider or None,
        router_correlation_id=router_correlation_id,
    )


def _get_spend_logs_metadata(
    metadata: dict | None,
    applied_guardrails: list[str] | None = None,
    batch_models: list[str] | None = None,
    batch_successful_requests: int | None = None,
    batch_failed_requests: int | None = None,
    mcp_tool_call_metadata: StandardLoggingMCPToolCall | None = None,
    vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None = None,
    guardrail_information: list[StandardLoggingGuardrailInformation] | None = None,
    usage_object: dict | None = None,
    model_map_information: StandardLoggingModelInformation | None = None,
    cold_storage_object_key: str | None = None,
    litellm_overhead_time_ms: float | None = None,
    cost_breakdown: CostBreakdown | None = None,
    litellm_call_id: str | None = None,
    autorouter_savings: float | None = None,
    router_metadata: SpendLogsRouterMetadata | None = None,
) -> SpendLogsMetadata:
    if metadata is None:
        return SpendLogsMetadata(
            user_api_key=None,
            user_api_key_alias=None,
            user_api_key_team_id=None,
            user_api_key_project_id=None,
            user_api_key_project_alias=None,
            user_api_key_org_id=None,
            user_api_key_user_id=None,
            user_api_key_team_alias=None,
            spend_logs_metadata=None,
            requester_ip_address=None,
            additional_usage_values=None,
            applied_guardrails=None,
            status=None or "success",
            error_information=None,
            proxy_server_request=None,
            batch_models=None,
            batch_successful_requests=None,
            batch_failed_requests=None,
            mcp_tool_call_metadata=None,
            vector_store_request_metadata=None,
            model_map_information=None,
            usage_object=None,
            guardrail_information=None,
            internal_call_origin=None,
            eval_information=None,
            cold_storage_object_key=cold_storage_object_key,
            litellm_overhead_time_ms=None,
            attempted_retries=None,
            max_retries=None,
            attempted_fallbacks=None,
            original_model_group=None,
            cost_breakdown=None,
            compression_savings=None,
            autorouter_savings=autorouter_savings,
            litellm_gateway_injected_cache=None,
            litellm_call_id=litellm_call_id,
            router_metadata=router_metadata,
        )
    verbose_proxy_logger.debug(
        "getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys()))
    )

    # Filter the metadata dictionary to include only the specified keys
    clean_metadata: Final = SpendLogsMetadata(
        **{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key != "router_metadata"},
        router_metadata=router_metadata,
    )
    _raw_key: Final = clean_metadata.get("user_api_key")
    _trusted_hash: Final = metadata.get("user_api_key_hash")
    _already_redacted: Final = (
        isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == _raw_key
    )
    clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_redacted=_already_redacted)
    clean_metadata["applied_guardrails"] = applied_guardrails
    clean_metadata["batch_models"] = batch_models
    clean_metadata["batch_successful_requests"] = batch_successful_requests
    clean_metadata["batch_failed_requests"] = batch_failed_requests
    clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata
    clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload(
        vector_store_request_metadata
    )
    clean_metadata["guardrail_information"] = _sanitize_guardrail_information_for_spend_logs(guardrail_information)
    clean_metadata["usage_object"] = usage_object
    clean_metadata["model_map_information"] = model_map_information
    clean_metadata["cold_storage_object_key"] = cold_storage_object_key
    clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms
    clean_metadata["cost_breakdown"] = cost_breakdown
    clean_metadata["autorouter_savings"] = autorouter_savings
    clean_metadata["litellm_call_id"] = litellm_call_id

    return clean_metadata


BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost"


def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | None:
    standard_logging_payload = kwargs.get("standard_logging_object")
    candidate_ids: Final = (
        response_obj.get("id"),
        standard_logging_payload.get("id") if isinstance(standard_logging_payload, dict) else None,
        kwargs.get("litellm_call_id"),
    )
    resolved_id: Final = next(
        (candidate for candidate in candidate_ids if isinstance(candidate, str) and candidate), None
    )
    if resolved_id is not None and call_type == CallTypes.aretrieve_batch.value:
        return f"{resolved_id}{BATCH_COST_REQUEST_ID_SUFFIX}"
    return resolved_id


_MISSING_ATTRIBUTE: Final = object()


def _attribute_or_missing(source: object, name: str) -> object:
    return getattr(source, name, _MISSING_ATTRIBUTE)


@runtime_checkable
class _ModelDumpable(Protocol):
    def model_dump(self) -> object: ...


def _dumped_usage_info(usage_info: object) -> object:
    if isinstance(usage_info, _ModelDumpable):
        return usage_info.model_dump()
    instance_dict: Final = _attribute_or_missing(usage_info, "__dict__")
    if instance_dict is not _MISSING_ATTRIBUTE:
        return instance_dict
    return usage_info


def _extract_usage_for_ocr_call(response_obj: object, response_obj_dict: dict) -> dict:
    """
    Extract usage information for OCR/AOCR calls.

    OCR responses use usage_info (with pages_processed) instead of token-based usage.

    Args:
        response_obj: The raw response object (can be dict, BaseModel, or other)
        response_obj_dict: Dictionary representation of the response object

    Returns:
        A dict with prompt_tokens=0, completion_tokens=0, total_tokens=0,
        and pages_processed from usage_info.
    """
    usage_info = None

    # Try to extract usage_info from dict
    if isinstance(response_obj_dict, dict) and "usage_info" in response_obj_dict:
        usage_info = response_obj_dict.get("usage_info")

    # Try to extract usage_info from object attributes if not found in dict
    if not usage_info:
        attribute_usage_info: Final = _attribute_or_missing(response_obj, "usage_info")
        if attribute_usage_info is not _MISSING_ATTRIBUTE:
            usage_info = _dumped_usage_info(attribute_usage_info)

    # For OCR, we track pages instead of tokens
    if usage_info is not None:
        # Handle dict or object with attributes
        if isinstance(usage_info, dict):
            result: Final = {
                "prompt_tokens": 0,  # OCR doesn't use traditional tokens
                "completion_tokens": 0,
                "total_tokens": 0,
            }
            # Add all fields from usage_info, including pages_processed
            for key, value in usage_info.items():
                result[key] = value
            # Ensure pages_processed exists
            if "pages_processed" not in result:
                result["pages_processed"] = 0
            return result
        else:
            return {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
                "pages_processed": 0,
            }
    else:
        return {}


def get_request_model_access_groups(kwargs: Mapping[str, object] | None) -> tuple[str, ...]:
    """Model access groups that authorized this request, as stamped onto request metadata at auth time."""
    if kwargs is None:
        return ()

    standard_logging_payload: Final = kwargs.get("standard_logging_object")
    if isinstance(standard_logging_payload, Mapping):
        from_payload: Final = coerce_model_access_groups(standard_logging_payload.get("request_model_access_groups"))
        if from_payload:
            return from_payload

    litellm_params: Final = kwargs.get("litellm_params")
    if not isinstance(litellm_params, Mapping):
        return ()
    return request_model_access_groups_from_litellm_params(litellm_params)


def _sl_attribution_fallback(
    standard_logging_payload: StandardLoggingPayload | None,
    field: Literal["model_id", "model_group", "api_base", "custom_llm_provider"],
) -> str:
    if standard_logging_payload is None:
        return ""
    return standard_logging_payload.get(field) or ""


def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogsPayload:
    if kwargs is None:
        kwargs = {}

    if response_obj is None:
        response_obj = {}
    elif not isinstance(response_obj, BaseModel) and not isinstance(response_obj, dict):
        response_obj = {"result": str(response_obj)}
    # standardize this function to be used across, s3, dynamoDB, langfuse logging
    litellm_params: Final = kwargs.get("litellm_params", {})
    metadata: Final = get_litellm_metadata_from_kwargs(kwargs)
    completion_start_time: Final = kwargs.get("completion_start_time", end_time)
    call_type: Final = kwargs.get("call_type")
    cache_hit: Final = kwargs.get("cache_hit", False)

    # Convert response_obj to dict first
    if isinstance(response_obj, dict):
        response_obj_dict = response_obj
    elif isinstance(response_obj, BaseModel):
        response_obj_dict = response_obj.model_dump()
    else:
        response_obj_dict = {}

    # Handle OCR responses which use usage_info instead of usage
    usage: dict = {}
    if call_type in ["ocr", "aocr"]:
        usage = _extract_usage_for_ocr_call(response_obj, response_obj_dict)
    elif not is_unbilled_non_inference_call(call_type, metadata, response_obj_dict):
        # Use response_obj_dict instead of response_obj to avoid calling .get() on Pydantic models
        _usage: Final = response_obj_dict.get("usage", None) or {}
        if isinstance(_usage, litellm.Usage):
            usage = dict(_usage)
        elif isinstance(_usage, dict):
            usage = _usage

    # A request that failed mid-stream has no usable response_obj usage, but the
    # streaming handler may have recovered the usage from the chunks already
    # delivered. Honor that override so the partial usage lands in spend tracking.
    _combined_usage: Final = kwargs.get("combined_usage_object")
    if not usage and isinstance(_combined_usage, litellm.Usage):
        usage = _combined_usage.model_dump()

    id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs)
    standard_logging_payload: Final = cast(StandardLoggingPayload | None, kwargs.get("standard_logging_object", None))

    end_user_id = get_end_user_id_for_cost_tracking(litellm_params)

    api_key = metadata.get("user_api_key", "")

    standard_logging_prompt_tokens: int = 0
    standard_logging_completion_tokens: int = 0
    standard_logging_total_tokens: int = 0
    if standard_logging_payload is not None:
        standard_logging_prompt_tokens = standard_logging_payload.get("prompt_tokens", 0)
        standard_logging_completion_tokens = standard_logging_payload.get("completion_tokens", 0)
        standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0)
    _trusted_hash = metadata.get("user_api_key_hash")
    _key_already_redacted = (
        isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == api_key
    )
    api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted) or ""

    if (
        standard_logging_payload is not None
    ):  # [TODO] migrate completely to sl payload. currently missing pass-through endpoint data
        api_key = (
            api_key
            or _redact_logged_api_key(
                standard_logging_payload["metadata"].get("user_api_key_hash"), already_redacted=True
            )
            or ""
        )
        end_user_id = end_user_id or standard_logging_payload["metadata"].get("user_api_key_end_user_id")
    request_tags = safe_dumps(metadata.get("tags", [])) if isinstance(metadata.get("tags", []), list) else "[]"
    if (
        standard_logging_payload is not None and standard_logging_payload.get("request_tags") is not None
    ):  # use 'tags' from standard logging payload instead
        request_tags = safe_dumps(standard_logging_payload["request_tags"])

    _model_id: Final = metadata.get("model_info", {}).get("id", "") or _sl_attribution_fallback(
        standard_logging_payload, "model_id"
    )
    _model_group: Final = metadata.get("model_group", "") or _sl_attribution_fallback(
        standard_logging_payload, "model_group"
    )
    _api_base: Final = litellm_params.get("api_base", "") or _sl_attribution_fallback(
        standard_logging_payload, "api_base"
    )

    # Extract overhead from hidden_params if available
    litellm_overhead_time_ms = None
    if standard_logging_payload is not None:
        hidden_params: Final = standard_logging_payload.get("hidden_params", {})
        litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms")

    custom_llm_provider: Final = (
        kwargs.get("custom_llm_provider")
        or _sl_attribution_fallback(standard_logging_payload, "custom_llm_provider")
        or None
    )
    raw_model: Final = cast(str, kwargs.get("model") or "")
    model_name: Final = (
        standard_logging_payload.get("model") if standard_logging_payload is not None else None
    ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
    litellm_call_id: Final = cast(
        str | None,
        kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"),
    )

    # clean up litellm metadata
    clean_metadata = _get_spend_logs_metadata(
        metadata,
        applied_guardrails=(
            standard_logging_payload["metadata"].get("applied_guardrails", None)
            if standard_logging_payload is not None
            else None
        ),
        batch_models=(
            standard_logging_payload.get("hidden_params", {}).get("batch_models", None)
            if standard_logging_payload is not None
            else None
        ),
        batch_successful_requests=(
            standard_logging_payload.get("hidden_params", {}).get("batch_successful_requests", None)
            if standard_logging_payload is not None
            else None
        ),
        batch_failed_requests=(
            standard_logging_payload.get("hidden_params", {}).get("batch_failed_requests", None)
            if standard_logging_payload is not None
            else None
        ),
        mcp_tool_call_metadata=(
            standard_logging_payload["metadata"].get("mcp_tool_call_metadata", None)
            if standard_logging_payload is not None
            else None
        ),
        vector_store_request_metadata=(
            standard_logging_payload["metadata"].get("vector_store_request_metadata", None)
            if standard_logging_payload is not None
            else None
        ),
        usage_object=(
            standard_logging_payload["metadata"].get("usage_object", None)
            if standard_logging_payload is not None
            else None
        ),
        model_map_information=(
            standard_logging_payload["model_map_information"] if standard_logging_payload is not None else None
        ),
        guardrail_information=(
            standard_logging_payload.get("guardrail_information", None)
            if standard_logging_payload is not None
            else (metadata.get("standard_logging_guardrail_information", None) if metadata is not None else None)
        ),
        cold_storage_object_key=(
            standard_logging_payload["metadata"].get("cold_storage_object_key", None)
            if standard_logging_payload is not None
            else None
        ),
        litellm_overhead_time_ms=litellm_overhead_time_ms,
        cost_breakdown=(
            standard_logging_payload.get("cost_breakdown", None) if standard_logging_payload is not None else None
        ),
        autorouter_savings=(
            standard_logging_payload.get("autorouter_savings", None) if standard_logging_payload is not None else None
        ),
        litellm_call_id=litellm_call_id,
        router_metadata=_get_router_metadata_for_spend_log(
            metadata=metadata,
            requested_model=_model_group,
            selected_model=model_name,
            selected_provider=custom_llm_provider,
            router_correlation_id=litellm_call_id,
        ),
    )

    special_usage_fields: Final = ["completion_tokens", "prompt_tokens", "total_tokens"]
    additional_usage_values = {}
    for k, v in usage.items():
        if k not in special_usage_fields:
            if isinstance(v, BaseModel):
                v = v.model_dump()
            additional_usage_values.update({k: v})
    prompt_tokens_details = additional_usage_values.get("prompt_tokens_details")
    if not isinstance(prompt_tokens_details, dict):
        usage_object: Final = clean_metadata.get("usage_object")
        if isinstance(usage_object, dict):
            prompt_tokens_details = usage_object.get("prompt_tokens_details")
    if isinstance(prompt_tokens_details, dict):
        if "cache_read_input_tokens" not in additional_usage_values:
            cached_tokens: Final = prompt_tokens_details.get("cached_tokens")
            if isinstance(cached_tokens, int) and cached_tokens > 0:
                additional_usage_values["cache_read_input_tokens"] = cached_tokens
        if "cache_creation_input_tokens" not in additional_usage_values:
            cache_write_tokens: Final = prompt_tokens_details.get("cache_write_tokens") or prompt_tokens_details.get(
                "cache_creation_tokens"
            )
            if isinstance(cache_write_tokens, int) and cache_write_tokens > 0:
                additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens
    clean_metadata["additional_usage_values"] = additional_usage_values

    if litellm.cache is not None:
        cache_key = litellm.cache.get_cache_key(**kwargs)
    else:
        cache_key = "Cache OFF"
    if cache_hit is True:
        import time

        id = f"{id}_cache_hit{time.time()}"  # SpendLogs does not allow duplicate request_id

    mcp_namespaced_tool_name = None
    mcp_tool_call_metadata: Final[StandardLoggingMCPToolCall | None] = clean_metadata.get("mcp_tool_call_metadata")
    if mcp_tool_call_metadata is not None:
        mcp_namespaced_tool_name = mcp_tool_call_metadata.get("namespaced_tool_name", None)

    # Extract agent_id for A2A requests (set directly on model_call_details)
    agent_id: Final[str | None] = kwargs.get("agent_id") or metadata.get("agent_id")

    try:
        payload: Final[SpendLogsPayload] = SpendLogsPayload(
            request_id=str(id),
            call_type=call_type or "",
            api_key=str(api_key),
            cache_hit=str(cache_hit),
            startTime=_ensure_datetime_utc(start_time),
            endTime=_ensure_datetime_utc(end_time),
            completionStartTime=_ensure_datetime_utc(completion_start_time),
            model=model_name,
            user=metadata.get("user_api_key_user_id", "") or "",
            team_id=metadata.get("user_api_key_team_id", "") or "",
            organization_id=metadata.get("user_api_key_org_id") or "",
            metadata=safe_dumps(clean_metadata),
            cache_key=cache_key,
            spend=kwargs.get("response_cost", 0),
            total_tokens=usage.get("total_tokens", standard_logging_total_tokens),
            prompt_tokens=usage.get("prompt_tokens", standard_logging_prompt_tokens),
            completion_tokens=usage.get("completion_tokens", standard_logging_completion_tokens),
            request_tags=request_tags,
            end_user=end_user_id or "",
            api_base=_api_base,
            model_group=_model_group,
            model_id=_model_id,
            mcp_namespaced_tool_name=mcp_namespaced_tool_name,
            agent_id=agent_id,
            requester_ip_address=clean_metadata.get("requester_ip_address", None),
            custom_llm_provider=custom_llm_provider or "",
            messages=_get_messages_for_spend_logs_payload(
                standard_logging_payload=standard_logging_payload, metadata=metadata
            ),
            response=_get_response_for_spend_logs_payload(payload=standard_logging_payload, kwargs=kwargs),
            proxy_server_request=_get_proxy_server_request_for_spend_logs_payload(
                metadata=metadata, litellm_params=litellm_params, kwargs=kwargs
            ),
            session_id=_get_session_id_for_spend_log(
                kwargs=kwargs,
                standard_logging_payload=standard_logging_payload,
            ),
            request_duration_ms=_get_request_duration_ms(start_time, end_time),
            status=_get_status_for_spend_log(
                metadata=metadata,
            ),
        )

        verbose_proxy_logger.debug(
            "SpendTable: created payload - request_id: %s, model: %s, spend: %s",
            payload.get("request_id"),
            payload.get("model"),
            payload.get("spend"),
        )

        # Explicitly clear large intermediate objects to reduce memory pressure
        del response_obj_dict, usage, clean_metadata, additional_usage_values

        return payload
    except Exception as e:
        spend_log_error("Error creating spendlogs object - %s", str(e), exc=e)
        raise e


def _get_session_id_for_spend_log(
    kwargs: dict,
    standard_logging_payload: StandardLoggingPayload | None,
) -> str:
    """
    Get the session id for the spend log.

    This ensures each spend log is associated with a unique session id.

    """
    from litellm._uuid import uuid

    if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None:
        return str(standard_logging_payload.get("trace_id"))

    # Users can dynamically set the trace_id for each request by passing `litellm_trace_id` in kwargs
    if kwargs.get("litellm_trace_id") is not None:
        return str(kwargs.get("litellm_trace_id"))

    # Ensure we always have a session id, if none is provided
    return str(uuid.uuid4())


def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> int | None:
    """Compute request duration in milliseconds from start and end times."""
    try:
        return int((end_time - start_time).total_seconds() * 1000)
    except Exception:
        return None


def _ensure_datetime_utc(timestamp: datetime) -> datetime:
    """Helper to ensure datetime is in UTC"""
    timestamp = timestamp.astimezone(timezone.utc)
    return timestamp


async def _query_raw_rows(
    prisma_client: PrismaClient,
    sql_query: str,
    *args: object,
) -> Sequence[Mapping[str, object]] | None:
    return await prisma_client.db.query_raw(sql_query, *args)


async def get_spend_by_team(
    start_date: dt,
    end_date: dt,
    team_id: str | None,
    prisma_client: PrismaClient,
):
    sql_query: Final = """
    WITH SpendByModelApiKey AS (
        SELECT
            date_trunc('day', sl."startTime") AS group_by_day,
            COALESCE(tt.team_alias, 'Unassigned Team') AS team_name,
            sl.model,
            sl.api_key,
            SUM(sl.spend) AS model_api_spend,
            SUM(sl.total_tokens) AS model_api_tokens
        FROM
            "LiteLLM_SpendLogs" sl
        LEFT JOIN
            "LiteLLM_TeamTable" tt
        ON
            sl.team_id = tt.team_id
        WHERE
            sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
            AND sl."startTime" <  (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
            AND ($3::text IS NULL OR sl.team_id = $3)
        GROUP BY
            date_trunc('day', sl."startTime"),
            tt.team_alias,
            sl.model,
            sl.api_key
    )
        SELECT
            group_by_day,
            jsonb_agg(jsonb_build_object(
                'team_name', team_name,
                'total_spend', total_spend,
                'metadata', metadata
            )) AS teams
        FROM (
            SELECT
                group_by_day,
                team_name,
                SUM(model_api_spend) AS total_spend,
                jsonb_agg(jsonb_build_object(
                    'model', model,
                    'api_key', api_key,
                    'spend', model_api_spend,
                    'total_tokens', model_api_tokens
                )) AS metadata
            FROM
                SpendByModelApiKey
            GROUP BY
                group_by_day,
                team_name
        ) AS aggregated
        GROUP BY
            group_by_day
        ORDER BY
            group_by_day;
    """

    db_response: Final = await _query_raw_rows(prisma_client, sql_query, start_date, end_date, team_id)
    if db_response is None:
        return []

    return db_response


async def get_spend_by_team_and_customer(
    start_date: dt,
    end_date: dt,
    team_id: str,
    customer_id: str,
    prisma_client: PrismaClient,
):
    sql_query: Final = """
    WITH SpendByModelApiKey AS (
        SELECT
            date_trunc('day', sl."startTime") AS group_by_day,
            COALESCE(tt.team_alias, 'Unassigned Team') AS team_name,
            sl.end_user AS customer,
            sl.model,
            sl.api_key,
            SUM(sl.spend) AS model_api_spend,
            SUM(sl.total_tokens) AS model_api_tokens
        FROM 
            "LiteLLM_SpendLogs" sl
        LEFT JOIN 
            "LiteLLM_TeamTable" tt 
        ON 
            sl.team_id = tt.team_id
        WHERE
            sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
            AND sl."startTime" <  (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
            AND sl.team_id = $3
            AND sl.end_user = $4
        GROUP BY
            date_trunc('day', sl."startTime"),
            tt.team_alias,
            sl.end_user,
            sl.model,
            sl.api_key
    )
        SELECT
            group_by_day,
            jsonb_agg(jsonb_build_object(
                'team_name', team_name,
                'customer', customer,
                'total_spend', total_spend,
                'metadata', metadata
            )) AS teams_customers
        FROM (
            SELECT
                group_by_day,
                team_name,
                customer,
                SUM(model_api_spend) AS total_spend,
                jsonb_agg(jsonb_build_object(
                    'model', model,
                    'api_key', api_key,
                    'spend', model_api_spend,
                    'total_tokens', model_api_tokens
                )) AS metadata
            FROM 
                SpendByModelApiKey
            GROUP BY
                group_by_day,
                team_name,
                customer
        ) AS aggregated
        GROUP BY
            group_by_day
        ORDER BY
            group_by_day;
    """

    db_response: Final = await _query_raw_rows(prisma_client, sql_query, start_date, end_date, team_id, customer_id)
    if db_response is None:
        return []

    return db_response


def _get_messages_for_spend_logs_payload(
    standard_logging_payload: StandardLoggingPayload | None,
    metadata: dict | None = None,
) -> str:
    if _should_store_prompts_and_responses_in_spend_logs():
        if standard_logging_payload is not None:
            call_type: Final = standard_logging_payload.get("call_type", "")
            if call_type == "_arealtime":
                messages: Final = standard_logging_payload.get("messages")
                if messages is not None:
                    try:
                        return safe_dumps(messages)
                    except Exception:
                        return "{}"
    return "{}"


_SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields"})


def _sanitize_request_body_for_spend_logs_payload(
    request_body: dict,
    visited: set | None = None,
    max_string_length_prompt_in_db: int | None = None,
) -> dict:
    """
    Recursively sanitize request body to prevent logging large base64 strings or other large values.
    Truncates strings longer than MAX_STRING_LENGTH_PROMPT_IN_DB characters and handles nested dictionaries.

    Also strips keys listed in _SENSITIVE_REQUEST_BODY_KEYS (e.g. secret_fields
    which contains raw HTTP headers including Authorization tokens).
    """
    from litellm.constants import (
        LITELLM_TRUNCATED_PAYLOAD_FIELD,
        LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
    )

    if visited is None:
        visited = set()
    if max_string_length_prompt_in_db is None:
        max_string_length_prompt_in_db = _get_max_string_length_prompt_in_db()

    # Get the object's memory address to track visited objects
    obj_id: Final = id(request_body)
    if obj_id in visited:
        return {}
    visited.add(obj_id)

    def _sanitize_value(value: object) -> object:
        if isinstance(value, dict):
            return _sanitize_request_body_for_spend_logs_payload(value, visited, max_string_length_prompt_in_db)
        elif isinstance(value, list):
            return [_sanitize_value(item) for item in value]
        elif isinstance(value, str):
            if len(value) > max_string_length_prompt_in_db:
                # Keep 35% from beginning and 65% from end (end is usually more important)
                # This split ensures we keep more context from the end of conversations
                start_ratio: Final = 0.35
                end_ratio: Final = 0.65

                # Calculate character distribution
                start_chars: Final = int(max_string_length_prompt_in_db * start_ratio)
                end_chars = int(max_string_length_prompt_in_db * end_ratio)

                # Ensure we don't exceed the total limit
                total_keep: Final = start_chars + end_chars
                if total_keep > max_string_length_prompt_in_db:
                    end_chars = max_string_length_prompt_in_db - start_chars

                # If the string length is less than what we want to keep, just truncate normally
                if len(value) <= max_string_length_prompt_in_db:
                    return value

                # Calculate how many characters are being skipped
                skipped_chars: Final = len(value) - total_keep

                # Build the truncated string: beginning + truncation marker + end
                truncated_value: Final = (
                    f"{value[:start_chars]}"
                    f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. "
                    f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..."
                    f"{value[-end_chars:]}"
                )
                return truncated_value
            return value
        return value

    return {k: _sanitize_value(v) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS}


# Quoted-key form: ``"input"`` / ``'messages'`` / ``"prompt"`` followed by
# ``:``. Covers JSON bodies and Python dict-reprs in provider error strings.
# ``prompt`` is included for ``/v1/completions``-style payloads where the user
# input lives under a top-level ``prompt`` key rather than ``messages``.
_ERROR_MESSAGE_PROMPT_LEAK_KEYS: Final = ("input", "messages", "prompt")


# Assignment-style keys: Pydantic v2 validation errors render the offending
# value as ``input_value=<repr>`` inside ``[type=..., input_value=...,
# input_type=...]``. The same prompt body that would appear under an
# ``"input"`` JSON key is echoed here as a Python repr, so we redact it
# under the same store_prompts_in_spend_logs gate.
_ERROR_MESSAGE_ASSIGN_LEAK_KEYS: Final = ("input_value",)


_SENSITIVE_KEY_START_PATTERN: Final = re.compile(
    r"(?:"
    r"['\"](?:" + "|".join(_ERROR_MESSAGE_PROMPT_LEAK_KEYS) + r")['\"]\s*:\s*"
    r"|"
    r"\b(?:" + "|".join(_ERROR_MESSAGE_ASSIGN_LEAK_KEYS) + r")\s*=\s*"
    r")"
)


def _scan_quoted_string_end(text: str, start: int, quote: str) -> int:
    """
    Given ``text[start] == quote`` (``'`` or ``"``), return the index just
    past the matching close quote, honoring backslash escapes. Returns
    ``-1`` if unterminated.
    """
    n: Final = len(text)
    i = start + 1
    while i < n:
        c = text[i]
        if c == "\\":
            i += 2
            continue
        if c == quote:
            return i + 1
        i += 1
    return -1


def _scan_balanced_value_end(text: str, start: int) -> int:
    """
    Given ``text[start]`` is ``[``, ``{``, ``'`` or ``"``, return the index
    just past the matching close, accounting for nested brackets and
    quoted strings (with escape sequences). Returns ``-1`` if the
    structure is unterminated.

    Implemented iteratively (no self-recursion): the bracket scanner
    inlines a quote-skip helper rather than re-entering itself, since
    JSON-style values cannot contain another bracket *as a first char*
    inside a quoted string — only the quote-skip case can occur.
    """
    n: Final = len(text)
    if start >= n:
        return -1
    first: Final = text[start]
    if first in ("'", '"'):
        return _scan_quoted_string_end(text, start, first)
    if first == "[":
        close = "]"
    elif first == "{":
        close = "}"
    else:
        return -1
    depth = 0
    i = start
    while i < n:
        c = text[i]
        if c in ("'", '"'):
            end = _scan_quoted_string_end(text, i, c)
            if end == -1:
                return -1
            i = end
            continue
        if c == first:
            depth += 1
        elif c == close:
            depth -= 1
            if depth == 0:
                return i + 1
        i += 1
    return -1


def _redact_prompt_leaks_in_error_string(text: str) -> str:
    """
    Strip echoed request input from provider error strings.

    Provider validation errors (e.g. OpenAI ``RateLimitError`` carrying 178
    pydantic validation errors, each with its own ``'input': [...]`` field)
    embed the full request body in their message. When prompts must not be
    stored in spend logs, that echo is a back-door leak.

    Two leak shapes are handled:

    - Quoted-key form — ``"<key>": <value>`` where ``key`` is ``input``,
      ``messages`` or ``prompt`` (covers JSON bodies, Python dict-reprs,
      and ``/v1/completions`` payloads).
    - Assignment form — ``input_value=<value>`` from Pydantic v2 validation
      errors, which render the offending value as a Python repr inside
      ``[type=..., input_value=..., input_type=...]``.

    The value scan understands nested ``[]`` / ``{}`` and quoted strings,
    so multi-modal payloads (``'messages': [{'content': [{...}]}]``) and
    user text containing brackets (``"secret[123"``) are handled correctly.
    """
    if not text:
        return text
    redaction: Final = f'"{REDACTED_BY_LITELM_STRING}"'
    out: Final[list[str]] = []
    n: Final = len(text)
    pos = 0
    while pos < n:
        m = _SENSITIVE_KEY_START_PATTERN.search(text, pos)
        if not m:
            out.append(text[pos:])
            break
        out.append(text[pos : m.end()])
        v_start = m.end()
        if v_start >= n:
            break
        first = text[v_start]
        if first in ("[", "{", "'", '"'):
            v_end = _scan_balanced_value_end(text, v_start)
            if v_end == -1:
                # Unterminated value — redact through the rest of the string
                # so a malformed leak can't slip past.
                out.append(redaction)
                pos = n
                break
            out.append(redaction)
            pos = v_end
        else:
            # Unquoted scalar (number, null, bare identifier) — not a leak
            # carrier, leave intact and resume after the key match.
            pos = v_start
    return "".join(out)


def _sanitize_guardrail_information_for_spend_logs(
    guardrail_information: list[StandardLoggingGuardrailInformation] | None,
) -> list[StandardLoggingGuardrailInformation] | None:
    """
    When ``store_prompts_in_spend_logs`` is False, redact prompt-carrying fields
    (``guardrail_request``, ``guardrail_response``, ``match_details``,
    ``classification``) before they land in ``LiteLLM_SpendLogs.metadata``.

    Guardrail hooks may echo the LLM request payload back into
    ``guardrail_response``, and two first-party hooks
    (``block_code_execution``, ``litellm_content_filter``) inline user-prompt
    substrings into ``match_details`` / ``classification`` too, so the flag
    must cover all four fields. Every other typed field on the entry (name,
    provider, mode, status, timings, action, violation_categories, risk_score,
    masked_entity_count, ...) is preserved so guardrail dashboards keep
    working.

    ``guardrail_information`` is typed ``Optional[List[...]]`` but at least
    one writer (``xecguard``) assigns a bare dict, so normalize to a list
    here to match OTEL's defensive read pattern; otherwise iteration would
    yield the dict's keys and crash the whole spend-log write.
    """
    if guardrail_information is None or _should_store_prompts_and_responses_in_spend_logs():
        return guardrail_information
    entries: Final = [guardrail_information] if isinstance(guardrail_information, dict) else guardrail_information
    return [_redact_prompt_fields_in_guardrail_entry(entry) for entry in entries if isinstance(entry, dict)]


_PROMPT_CARRYING_GUARDRAIL_FIELDS: Final = (
    "guardrail_request",
    "guardrail_response",
    "match_details",
    "classification",
)

_NUMERIC_COMPRESSION_STAT_KEYS: Final = (
    "tokens_before",
    "tokens_after",
    "tokens_saved",
    "compression_ratio",
)


def _numeric_compression_stats_from_guardrail_response(
    guardrail_response: object,
) -> dict[str, int | float] | None:
    if not isinstance(guardrail_response, dict):
        return None
    stats: Final = {
        key: value
        for key, value in guardrail_response.items()
        if key in _NUMERIC_COMPRESSION_STAT_KEYS and isinstance(value, (int, float)) and not isinstance(value, bool)
    }
    return stats or None


def _redact_prompt_fields_in_guardrail_entry(
    entry: StandardLoggingGuardrailInformation,
) -> StandardLoggingGuardrailInformation:
    """
    Replace prompt-carrying fields with the redaction marker. Purely numeric
    compression stats inside ``guardrail_response`` (e.g. Headroom's
    ``tokens_saved``) cannot carry prompt content, so they are preserved as a
    stats-only dict; spend aggregation reads them via
    ``extract_compression_saved_tokens``.
    """
    preserved_stats: Final = _numeric_compression_stats_from_guardrail_response(entry.get("guardrail_response"))
    redacted: Final[StandardLoggingGuardrailInformation] = {
        **entry,
        **{key: REDACTED_BY_LITELM_STRING for key in _PROMPT_CARRYING_GUARDRAIL_FIELDS if key in entry},
    }
    if preserved_stats is None:
        return redacted
    return {**redacted, "guardrail_response": preserved_stats}


def _sanitize_error_information_for_spend_logs(
    error_information: StandardLoggingPayloadErrorInformation | None,
) -> StandardLoggingPayloadErrorInformation | None:
    """
    Sanitize ``error_information`` before it lands in ``LiteLLM_SpendLogs.metadata``.

    Provider errors are stored verbatim via ``str(original_exception)``; those
    strings can echo the full request body, producing multi-megabyte spend-log
    rows.

    - Always: cap ``error_message`` and ``traceback`` with the existing
      ``MAX_STRING_LENGTH_PROMPT_IN_DB`` DB-storage safeguard.
    - When ``store_prompts_in_spend_logs`` is False: additionally redact
      ``'input'`` / ``'messages'`` / ``'prompt'`` values *and* Pydantic v2
      ``input_value=...`` assignments inside both ``error_message`` and
      ``traceback`` so prompts cannot leak through either field.

    Scoped to the spend-log path — OTEL/Datadog/etc. callbacks still receive
    the untruncated error per ``LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE``.
    """
    if error_information is None:
        return None

    sanitized = cast(dict, {**error_information})

    if not _should_store_prompts_and_responses_in_spend_logs():
        for field in ("error_message", "traceback"):
            value = sanitized.get(field)
            if isinstance(value, str):
                sanitized[field] = _redact_prompt_leaks_in_error_string(value)

    sanitized = _sanitize_request_body_for_spend_logs_payload(sanitized)
    return cast(StandardLoggingPayloadErrorInformation, sanitized)


def _convert_to_json_serializable_dict(obj: object, visited: set[int] | None = None, max_depth: int = 20) -> object:
    """
    Convert object to JSON-serializable dict, handling Pydantic models safely.

    This avoids pickle-based deepcopy which fails on Pydantic v2 models
    containing _thread.RLock objects.

    Args:
        obj: Object to convert (dict, list, Pydantic model, or primitive)
        visited: Set of object IDs to track circular references
        max_depth: Maximum recursion depth to prevent infinite recursion

    Returns:
        JSON-serializable version of the object
    """
    if max_depth <= 0:
        # Return a placeholder if max depth is exceeded
        return "<max_depth_exceeded>"

    if visited is None:
        visited = set()

    # Get the object's memory address to track visited objects
    obj_id: Final = id(obj)
    if obj_id in visited:
        # Circular reference detected, return placeholder
        return "<circular_reference>"

    # Only track mutable objects (dict, list, objects with __dict__)
    if isinstance(obj, (dict, list)) or hasattr(obj, "__dict__"):
        visited.add(obj_id)

    try:
        if isinstance(obj, BaseModel):
            # Use Pydantic's model_dump() instead of pickle
            result: Final = obj.model_dump()
            # Recursively process the dumped dict
            return _convert_to_json_serializable_dict(result, visited, max_depth - 1)
        elif isinstance(obj, dict):
            return {k: _convert_to_json_serializable_dict(v, visited, max_depth - 1) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [_convert_to_json_serializable_dict(item, visited, max_depth - 1) for item in obj]
        elif hasattr(obj, "__dict__"):
            # Handle objects with __dict__ attribute
            return _convert_to_json_serializable_dict(obj.__dict__, visited, max_depth - 1)
        else:
            # Primitives (str, int, float, bool, None) pass through
            return obj
    finally:
        # Remove from visited set when done processing this object
        if obj_id in visited:
            visited.remove(obj_id)


def _convert_mapping_to_json_serializable(obj: Mapping[str, object]) -> dict[str, object]:
    converted: Final = _convert_to_json_serializable_dict(obj)
    if isinstance(converted, dict):
        return converted
    return dict(obj)


def _get_proxy_server_request_for_spend_logs_payload(
    metadata: dict,
    litellm_params: dict,
    kwargs: dict | None = None,
) -> str:
    """
    Only store if _should_store_prompts_and_responses_in_spend_logs() is True

    If turn_off_message_logging is enabled, redact messages in the request body.
    """
    if _should_store_prompts_and_responses_in_spend_logs():
        _proxy_server_request: Final = cast(dict | None, litellm_params.get("proxy_server_request", {}))
        if _proxy_server_request is not None:
            _request_body = _proxy_server_request.get("body", {}) or {}

            if kwargs is not None:
                realtime_tools: Final = kwargs.get("realtime_tools")
                if realtime_tools:
                    _request_body = dict(_request_body)
                    _request_body["tools"] = realtime_tools

            # Apply message redaction if turn_off_message_logging is enabled
            if kwargs is not None:
                from litellm.litellm_core_utils.redact_messages import (
                    perform_redaction,
                    should_redact_message_logging,
                )

                # Build model_call_details dict to check redaction settings
                model_call_details: Final = {
                    "litellm_params": litellm_params,
                    "standard_callback_dynamic_params": kwargs.get("standard_callback_dynamic_params"),
                }

                # If redaction is enabled, convert to serializable dict before redacting
                if should_redact_message_logging(model_call_details=model_call_details):
                    _request_body = _convert_mapping_to_json_serializable(_request_body)
                    perform_redaction(model_call_details=_request_body, result=None)

            _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body)
            _request_body_json_str: Final = safe_dumps(_request_body)
            if LITELLM_TRUNCATED_PAYLOAD_FIELD in _request_body_json_str:
                verbose_proxy_logger.info(
                    "Spend Log: request body was truncated before storing in DB. %s",
                    LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
                )
            return _request_body_json_str
    return "{}"


def _get_vector_store_request_for_spend_logs_payload(
    vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None,
) -> list[StandardLoggingVectorStoreRequest] | None:
    """
    If user does not want to store prompts and responses, then remove the content from the vector store request metadata
    """
    if _should_store_prompts_and_responses_in_spend_logs():
        return vector_store_request_metadata

    # if user does not want to store prompts and responses, then remove the content from the vector store request metadata
    if vector_store_request_metadata is None:
        return None
    for vector_store_request in vector_store_request_metadata:
        vector_store_search_response: VectorStoreSearchResponse = (
            vector_store_request.get("vector_store_search_response") or VectorStoreSearchResponse()
        )
        response_data = vector_store_search_response.get("data", []) or []
        for response_item in response_data:
            for content_item in response_item.get("content", []) or []:
                if "text" in content_item:
                    content_item["text"] = REDACTED_BY_LITELM_STRING
    return vector_store_request_metadata


def _get_response_for_spend_logs_payload(
    payload: StandardLoggingPayload | None,
    kwargs: dict | None = None,
) -> str:
    if payload is None:
        return "{}"
    if _should_store_prompts_and_responses_in_spend_logs():
        response_obj: object = payload.get("response")
        if response_obj is None:
            return "{}"

        if kwargs is not None:
            realtime_tool_calls: Final = kwargs.get("realtime_tool_calls")
            if realtime_tool_calls and isinstance(response_obj, dict):
                response_obj = dict(response_obj)
                response_obj["tool_calls"] = realtime_tool_calls

        # Apply message redaction if turn_off_message_logging is enabled
        if kwargs is not None:
            from litellm.litellm_core_utils.redact_messages import (
                perform_redaction,
                should_redact_message_logging,
            )

            litellm_params: Final = kwargs.get("litellm_params", {})
            model_call_details: Final = {
                "litellm_params": litellm_params,
                "standard_callback_dynamic_params": kwargs.get("standard_callback_dynamic_params"),
            }

            # If redaction is enabled, convert to serializable dict before redacting
            if should_redact_message_logging(model_call_details=model_call_details):
                response_obj = _convert_to_json_serializable_dict(response_obj)
                response_obj = perform_redaction(model_call_details={}, result=response_obj)

        sanitized_wrapper: Final = _sanitize_request_body_for_spend_logs_payload({"response": response_obj})

        sanitized_response: Final = sanitized_wrapper.get("response", response_obj)

        if sanitized_response is None:
            return "{}"
        if isinstance(sanitized_response, str):
            result_str = strip_null_bytes(sanitized_response)
        else:
            result_str = safe_dumps(sanitized_response)
        if LITELLM_TRUNCATED_PAYLOAD_FIELD in result_str:
            verbose_proxy_logger.info(
                "Spend Log: response was truncated before storing in DB. %s",
                LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
            )
        return result_str
    return "{}"


def _should_store_prompts_and_responses_in_spend_logs() -> bool:
    from litellm.proxy.proxy_server import general_settings
    from litellm.secret_managers.main import get_secret_bool

    # Check general_settings (from DB or proxy_config.yaml)
    store_prompts_value: Final = general_settings.get("store_prompts_in_spend_logs")

    # Normalize case: handle True/true/TRUE, False/false/FALSE, None/null
    if store_prompts_value is True:
        return True
    elif isinstance(store_prompts_value, str):
        # Case-insensitive string comparison
        if store_prompts_value.lower() == "true":
            return True

    # Also check environment variable
    return get_secret_bool("STORE_PROMPTS_IN_SPEND_LOGS") is True


def _get_status_for_spend_log(
    metadata: dict,
) -> Literal["success", "failure"]:
    """
    Get the status for the spend log.

    It's only a failure if metadata.get("status") is "failure"
    """
    _status: Final[str | None] = metadata.get("status", None)
    if _status == "failure":
        return "failure"
    return "success"
