# Copyright 2026 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilities for Regional Access Boundary management."""

import asyncio
import copy
import datetime
import functools
import inspect
import logging
import os
import threading
from typing import NamedTuple, Optional, TYPE_CHECKING

from google.auth import _helpers
from google.auth import environment_vars

if TYPE_CHECKING:
    import google.auth.credentials
    import google.auth.transport

_LOGGER = logging.getLogger(__name__)


@functools.lru_cache()
def is_regional_access_boundary_enabled():
    """Checks if Regional Access Boundary is enabled via environment variable.

    The environment variable is interpreted as a boolean with the following
    (case-insensitive) rules:
    - "true", "1" are considered true.
    - Any other value (or unset) is considered false.

    Returns:
        bool: True if Regional Access Boundary is enabled, False otherwise.
    """
    value = os.environ.get(environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED)
    if value is None:
        return False

    return value.lower() in ("true", "1")


# The default lifetime for a cached Regional Access Boundary.
DEFAULT_REGIONAL_ACCESS_BOUNDARY_TTL = datetime.timedelta(hours=6)

# The period of time prior to the boundary's expiration when a background refresh
# is proactively triggered.
REGIONAL_ACCESS_BOUNDARY_REFRESH_THRESHOLD = datetime.timedelta(hours=1)

# The initial cooldown period for a failed Regional Access Boundary lookup.
DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN = datetime.timedelta(minutes=15)

# The maximum cooldown period for a failed Regional Access Boundary lookup.
MAX_REGIONAL_ACCESS_BOUNDARY_COOLDOWN = datetime.timedelta(hours=6)


# The header key used for Regional Access Boundaries.
_REGIONAL_ACCESS_BOUNDARY_HEADER = "x-allowed-locations"


class _RegionalAccessBoundaryData(NamedTuple):
    """Data container for a Regional Access Boundary snapshot.

    Attributes:
        encoded_locations (Optional[str]): The encoded Regional Access Boundary string.
        expiry (Optional[datetime.datetime]): The hard expiration time of the boundary data.
        cooldown_expiry (Optional[datetime.datetime]): The time until which further lookups are skipped.
        cooldown_duration (datetime.timedelta): The current duration for the exponential cooldown.
    """

    encoded_locations: Optional[str]
    expiry: Optional[datetime.datetime]
    cooldown_expiry: Optional[datetime.datetime]
    cooldown_duration: datetime.timedelta


class _RegionalAccessBoundaryManager(object):
    """Manages the Regional Access Boundary state and its background refresh.

    The actual data is held in an immutable `_RegionalAccessBoundaryData` object
    and is swapped atomically to ensure thread-safe, lock-free reads.
    """

    def __init__(self):
        self._data = _RegionalAccessBoundaryData(
            encoded_locations=None,
            expiry=None,
            cooldown_expiry=None,
            cooldown_duration=DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN,
        )
        self.refresh_manager = _RegionalAccessBoundaryRefreshManager()
        self._update_lock = threading.Lock()
        self._use_blocking_regional_access_boundary_lookup = False

    def __getstate__(self):
        """Pickle helper that serializes the _update_lock attribute."""
        state = self.__dict__.copy()
        state["_update_lock"] = None
        return state

    def __setstate__(self, state):
        """Pickle helper that deserializes the _update_lock attribute."""
        self.__dict__.update(state)
        self._update_lock = threading.Lock()

    def __eq__(self, other):
        """Checks if two managers are equal."""
        if not isinstance(other, _RegionalAccessBoundaryManager):
            return NotImplemented
        return (
            self._data == other._data
            and self._use_blocking_regional_access_boundary_lookup
            == other._use_blocking_regional_access_boundary_lookup
        )

    def enable_blocking_lookup(self):
        """Enables blocking Regional Access Boundary lookup.

        When enabled, the Regional Access Boundary lookup will be performed
        synchronously in the calling thread instead of asynchronously in a
        background thread.
        """
        self._use_blocking_regional_access_boundary_lookup = True

    def set_initial_regional_access_boundary(self, encoded_locations=None, expiry=None):
        """Manually sets the regional access boundary to the client provided initial values.

        Args:
            encoded_locations (Optional[str]): The encoded locations string.
            expiry (Optional[datetime.datetime]): The expiry time for the boundary.
                If encoded_locations is not provided, expiry is ignored.
        """
        if not encoded_locations:
            expiry = None

        self._data = _RegionalAccessBoundaryData(
            encoded_locations=encoded_locations,
            expiry=expiry,
            cooldown_expiry=None,
            cooldown_duration=DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN,
        )

    def apply_headers(self, headers):
        """Applies the Regional Access Boundary header to the provided dictionary.

        If the boundary is valid, the 'x-allowed-locations' header is added
        or updated. Otherwise, the header is removed to ensure no stale
        data is sent.

        Args:
            headers (MutableMapping[str, str]): The headers dictionary to update.
        """
        rab_data = self._data

        if rab_data.encoded_locations and (
            rab_data.expiry is not None and _helpers.utcnow() < rab_data.expiry
        ):
            headers[_REGIONAL_ACCESS_BOUNDARY_HEADER] = rab_data.encoded_locations
        else:
            headers.pop(_REGIONAL_ACCESS_BOUNDARY_HEADER, None)

    def _should_refresh(self):
        """Checks if the Regional Access Boundary data needs a refresh and is not in cooldown.

        Returns:
            bool: True if a refresh is required, False otherwise.
        """
        rab_data = self._data

        # Don't start a new refresh if the Regional Access Boundary info is still fresh.
        if (
            rab_data.encoded_locations
            and rab_data.expiry
            and _helpers.utcnow()
            < (rab_data.expiry - REGIONAL_ACCESS_BOUNDARY_REFRESH_THRESHOLD)
        ):
            return False

        # Don't start a new refresh if the cooldown is still in effect.
        if rab_data.cooldown_expiry and _helpers.utcnow() < rab_data.cooldown_expiry:
            return False

        return True

    def maybe_start_refresh(self, credentials, request):
        """Starts a background thread to refresh the Regional Access Boundary if needed.

        Args:
            credentials (google.auth.credentials.Credentials): The credentials to refresh.
            request (google.auth.transport.Request): The object used to make HTTP requests.
        """
        if not self._should_refresh():
            return

        # If all checks pass, start the background refresh.
        if self._use_blocking_regional_access_boundary_lookup:
            self.start_blocking_refresh(credentials, request)
        else:
            self.refresh_manager.start_refresh(credentials, request, self)

    async def maybe_start_refresh_async(self, credentials, request):
        """Starts a background refresh or performs a blocking refresh asynchronously.

        Args:
            credentials (google.auth.credentials.Credentials): The credentials to refresh.
            request (google.auth.aio.transport.Request): The object used to make HTTP requests.
        """
        if not self._should_refresh():
            return

        # If all checks pass, start the refresh.
        if self._use_blocking_regional_access_boundary_lookup:
            await self.start_blocking_refresh_async(credentials, request)
        else:
            self.refresh_manager.start_refresh(credentials, request, self)

    def start_blocking_refresh(self, credentials, request):
        """Initiates a blocking lookup of the Regional Access Boundary.

        If the lookup raises an exception, it is caught and logged as a warning,
        and the lookup is treated as a failure (entering cooldown). Exceptions
        are not propagated to the caller.

        Args:
            credentials (google.auth.credentials.Credentials): The credentials to refresh.
            request (google.auth.transport.Request): The object used to make HTTP requests.
        """
        # Async credentials do not support blocking lookups.
        if inspect.iscoroutinefunction(credentials._lookup_regional_access_boundary):
            if _helpers.is_logging_enabled(_LOGGER):
                _LOGGER.warning(
                    "Blocking Regional Access Boundary lookup is not supported for async credentials."
                )
            self.process_regional_access_boundary_info(None)
            return

        try:
            # The fail_fast parameter is set to True to ensure we don't block the calling
            # thread for too long. This will do two things: 1) set a timeout to 3s
            # instead of the default 120s and 2) ensure we do not retry at all
            regional_access_boundary_info = (
                credentials._lookup_regional_access_boundary(request, fail_fast=True)
            )
        except Exception as e:
            if _helpers.is_logging_enabled(_LOGGER):
                _LOGGER.warning(
                    "Blocking Regional Access Boundary lookup raised an exception: %s",
                    e,
                    exc_info=True,
                )
            regional_access_boundary_info = None

        self.process_regional_access_boundary_info(regional_access_boundary_info)

    async def start_blocking_refresh_async(self, credentials, request):
        """Initiates a blocking lookup of the Regional Access Boundary asynchronously.

        If the lookup raises an exception, it is caught and logged as a warning,
        and the lookup is treated as a failure (entering cooldown). Exceptions
        are not propagated to the caller.

        Args:
            credentials (google.auth.credentials.Credentials): The credentials to refresh.
            request (google.auth.aio.transport.Request): The object used to make HTTP requests.
        """
        try:
            # The fail_fast parameter is set to True to ensure we don't block the calling
            # thread for too long. This will do two things: 1) set a timeout to 3s
            # instead of the default 120s and 2) ensure we do not retry at all
            regional_access_boundary_info = (
                await credentials._lookup_regional_access_boundary(
                    request, fail_fast=True
                )
            )
        except Exception as e:
            if _helpers.is_logging_enabled(_LOGGER):
                _LOGGER.warning(
                    "Regional Access Boundary lookup raised an exception: %s",
                    e,
                    exc_info=True,
                )
            regional_access_boundary_info = None

        self.process_regional_access_boundary_info(regional_access_boundary_info)

    def process_regional_access_boundary_info(self, regional_access_boundary_info):
        """Processes the regional access boundary info and updates the state.

        Args:
            regional_access_boundary_info (Optional[Mapping[str, str]]): The regional access
                boundary info to process.
        """
        with self._update_lock:
            # Capture the current state before calculating updates.
            current_data = self._data

            if regional_access_boundary_info:
                # On success, update the boundary and its expiry, and clear any cooldown.
                encoded_locations = regional_access_boundary_info.get(
                    "encodedLocations"
                )
                updated_data = _RegionalAccessBoundaryData(
                    encoded_locations=encoded_locations,
                    expiry=_helpers.utcnow() + DEFAULT_REGIONAL_ACCESS_BOUNDARY_TTL,
                    cooldown_expiry=None,
                    cooldown_duration=DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN,
                )
                if _helpers.is_logging_enabled(_LOGGER):
                    _LOGGER.debug("Regional Access Boundary lookup successful.")
            else:
                # On failure, calculate cooldown and update state.
                if _helpers.is_logging_enabled(_LOGGER):
                    _LOGGER.warning(
                        "Regional Access Boundary lookup failed. Entering cooldown."
                    )

                next_cooldown_expiry = (
                    _helpers.utcnow() + current_data.cooldown_duration
                )
                next_cooldown_duration = min(
                    current_data.cooldown_duration * 2,
                    MAX_REGIONAL_ACCESS_BOUNDARY_COOLDOWN,
                )

                # If the refresh failed, we keep reusing the existing data unless
                # it has reached its hard expiration time.
                if current_data.expiry and _helpers.utcnow() > current_data.expiry:
                    next_encoded_locations = None
                    next_expiry = None
                else:
                    next_encoded_locations = current_data.encoded_locations
                    next_expiry = current_data.expiry

                updated_data = _RegionalAccessBoundaryData(
                    encoded_locations=next_encoded_locations,
                    expiry=next_expiry,
                    cooldown_expiry=next_cooldown_expiry,
                    cooldown_duration=next_cooldown_duration,
                )

            # Perform the atomic swap of the state object.
            self._data = updated_data


class _RegionalAccessBoundaryRefreshThread(threading.Thread):
    """Thread for background refreshing of the Regional Access Boundary."""

    def __init__(
        self,
        credentials: "google.auth.credentials.CredentialsWithRegionalAccessBoundary",  # noqa: F821
        request: "google.auth.transport.Request",  # noqa: F821
        rab_manager: "_RegionalAccessBoundaryManager",
    ):
        super().__init__()
        self.daemon = True
        self._credentials = credentials
        self._request = request
        self._rab_manager = rab_manager

    def run(self):
        """
        Performs the Regional Access Boundary lookup and updates the state.

        This method is run in a separate thread. It delegates the actual lookup
        to the credentials object's `_lookup_regional_access_boundary` method.
        Based on the lookup's outcome (success or complete failure after retries),
        it updates the cached Regional Access Boundary information,
        its expiry, its cooldown expiry, and its exponential cooldown duration.
        """
        # Catch exceptions (e.g., from the underlying transport) to prevent the
        # background thread from crashing. This ensures we can gracefully enter
        # an exponential cooldown state on failure.
        try:
            regional_access_boundary_info = (
                self._credentials._lookup_regional_access_boundary(self._request)
            )
        except Exception as e:
            if _helpers.is_logging_enabled(_LOGGER):
                _LOGGER.warning(
                    "Asynchronous Regional Access Boundary lookup raised an exception: %s",
                    e,
                    exc_info=True,
                )
            regional_access_boundary_info = None

        self._rab_manager.process_regional_access_boundary_info(
            regional_access_boundary_info
        )


class _RegionalAccessBoundaryRefreshManager(object):
    """Manages a thread for background refreshing of the Regional Access Boundary."""

    def __init__(self):
        self._lock = threading.Lock()
        self._worker = None

    def __getstate__(self):
        """Pickle helper that serializes the _lock and _worker attributes."""
        state = self.__dict__.copy()
        state["_lock"] = None
        state["_worker"] = None
        return state

    def __setstate__(self, state):
        """Pickle helper that deserializes the _lock and _worker attributes."""
        self.__dict__.update(state)
        self._lock = threading.Lock()
        self._worker = None

    def start_refresh(self, credentials, request, rab_manager):
        """
        Starts a background thread to refresh the Regional Access Boundary if one is not already running.

        Args:
            credentials (CredentialsWithRegionalAccessBoundary): The credentials
                to refresh.
            request (google.auth.transport.Request): The object used to make
                HTTP requests.
            rab_manager (_RegionalAccessBoundaryManager): The manager container to update.
        """
        with self._lock:
            if self._worker and self._worker.is_alive():
                # A refresh is already in progress.
                return

            try:
                copied_request = copy.deepcopy(request)
            except Exception as e:
                if _helpers.is_logging_enabled(_LOGGER):
                    _LOGGER.warning(
                        "Could not deepcopy transport for background RAB refresh. "
                        "Skipping background refresh to avoid thread safety issues. "
                        "Exception: %s",
                        e,
                    )
                return

            self._worker = _RegionalAccessBoundaryRefreshThread(
                credentials, copied_request, rab_manager
            )
            self._worker.start()


class _AsyncRegionalAccessBoundaryRefreshManager(object):
    """Manages a task for background refreshing of the Regional Access Boundary in async flows."""

    def __init__(self):
        self._lock = threading.Lock()
        self._worker_task = None

    def __getstate__(self):
        """Pickle helper that excludes the un-picklable _lock and _worker_task attributes from serialization."""
        state = self.__dict__.copy()
        state["_lock"] = None
        state["_worker_task"] = None
        return state

    def __setstate__(self, state):
        """Pickle helper that restores state and re-initializes the _lock and _worker_task attributes."""
        self.__dict__.update(state)
        self._lock = threading.Lock()
        self._worker_task = None

    def start_refresh(self, credentials, request, rab_manager):
        """
        Starts a background task to refresh the Regional Access Boundary if one is not already running.

        Args:
            credentials (CredentialsWithRegionalAccessBoundary): The credentials
                to refresh.
            request (google.auth.aio.transport.Request): The object used to make
                HTTP requests.
            rab_manager (_RegionalAccessBoundaryManager): The manager container to update.
        """
        with self._lock:
            if self._worker_task and not self._worker_task.done():
                # A refresh is already in progress.
                return

            async def _worker():
                try:
                    # credentials._lookup_regional_access_boundary should be async in the async creds class
                    regional_access_boundary_info = (
                        await credentials._lookup_regional_access_boundary(request)
                    )
                except Exception as e:
                    if _helpers.is_logging_enabled(_LOGGER):
                        _LOGGER.warning(
                            "Asynchronous Regional Access Boundary lookup raised an exception: %s",
                            e,
                            exc_info=True,
                        )
                    regional_access_boundary_info = None

                rab_manager.process_regional_access_boundary_info(
                    regional_access_boundary_info
                )

            coro = _worker()
            try:
                self._worker_task = asyncio.create_task(coro)
            except Exception:
                coro.close()
                raise


def _get_domain() -> str:
    """Dynamically determines the domain for IAM credentials based on active mTLS configuration.

    Returns:
        str: The dynamic domain string.
    """
    from google.auth.transport import _mtls_helper

    if (
        hasattr(_mtls_helper, "check_use_client_cert")
        and _mtls_helper.check_use_client_cert()
    ):
        return f"iamcredentials.mtls.{_helpers.DEFAULT_UNIVERSE_DOMAIN}"
    else:
        return f"iamcredentials.{_helpers.DEFAULT_UNIVERSE_DOMAIN}"


def get_service_account_rab_endpoint(service_account_email: str) -> str:
    """Builds the Regional Access Boundary lookup URL for service accounts.

    Args:
        service_account_email: The service account email.

    Returns:
        str: The complete lookup URL.
    """
    return f"https://{_get_domain()}/v1/projects/-/serviceAccounts/{service_account_email}/allowedLocations"


def get_workforce_pool_rab_endpoint(pool_id: str) -> str:
    """Builds the Regional Access Boundary lookup URL for workforce pools.

    Args:
        pool_id: The workforce pool ID.

    Returns:
        str: The complete lookup URL.
    """
    return f"https://{_get_domain()}/v1/locations/global/workforcePools/{pool_id}/allowedLocations"


def get_workload_identity_pool_rab_endpoint(project_number: str, pool_id: str) -> str:
    """Builds the Regional Access Boundary lookup URL for workload identity pools.

    Args:
        project_number: The Google Cloud project number.
        pool_id: The workload identity pool ID.

    Returns:
        str: The complete lookup URL.
    """
    return f"https://{_get_domain()}/v1/projects/{project_number}/locations/global/workloadIdentityPools/{pool_id}/allowedLocations"