"""Base client for gRPC services."""

from __future__ import annotations

from typing import TYPE_CHECKING, TypeVar

from radiens_core._transport._auth import AuthInterceptor
from radiens_core._transport._channel_pool import ChannelPool

if TYPE_CHECKING:
    from types import TracebackType

    from radiens_core.models.common import ServiceEndpoint

T = TypeVar("T", bound="BaseClient")


class BaseClient:
    """Base class for all Radiens gRPC clients.

    Handles channel pooling, authentication interceptors, and common
    lifecycle management.
    """

    def __init__(
        self,
        addresses: dict[str, str],
        auth_server_address: str,
        fallback_service: ServiceEndpoint,
    ) -> None:
        """Initialize the base client.

        Parameters
        ----------
        addresses
            Map of service name to "host:port" address.
        auth_server_address
            Address used by AuthInterceptor for server UUID fallback.
        fallback_service
            The default service endpoint to use if a requested service is
            not found in the address map.
        """
        self._addresses = addresses
        self._fallback_service = fallback_service
        self._pool = ChannelPool(
            interceptors=[AuthInterceptor(server_address=auth_server_address)],
        )

    # --- Lifecycle ---

    def close(self) -> None:
        """Close all gRPC channels managed by this client."""
        self._pool.close()

    def __enter__(self: T) -> T:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.close()

    # --- Internal helpers ---

    def _server_address(self, service: ServiceEndpoint) -> str:
        """Resolve the server address for a given service.

        Falls back to the configured ``fallback_service`` if the
        requested service name is not in the address map.
        """
        return self._addresses.get(service, self._addresses[self._fallback_service])
