from typing import Any, Optional, Union

from pydantic import BaseModel

from litellm.types.utils import HiddenParams


def add_retry_headers_to_response(
    response: Any,
    attempted_retries: int,
    max_retries: Optional[int] = None,
) -> Any:
    """
    Add retry headers to the request
    """

    if response is None or not isinstance(response, BaseModel):
        return response

    retry_headers = {
        "x-litellm-attempted-retries": attempted_retries,
    }
    if max_retries is not None:
        retry_headers["x-litellm-max-retries"] = max_retries

    hidden_params: Optional[Union[dict, HiddenParams]] = getattr(
        response, "_hidden_params", {}
    )

    if hidden_params is None:
        hidden_params = {}
    elif isinstance(hidden_params, HiddenParams):
        hidden_params = hidden_params.model_dump()

    hidden_params.setdefault("additional_headers", {})
    hidden_params["additional_headers"].update(retry_headers)

    setattr(response, "_hidden_params", hidden_params)

    return response