"""gRPC channel pooling with lifecycle management."""

from __future__ import annotations

import contextlib

import grpc

from radiens_core._constants import MAX_PB_MSG_SIZE_BYTES

_CHANNEL_OPTIONS: list[tuple[str, int | str | None]] = [
    ("grpc.max_receive_message_length", int(MAX_PB_MSG_SIZE_BYTES)),
    ("grpc.max_send_message_length", int(MAX_PB_MSG_SIZE_BYTES)),
]


class ChannelPool:
    """Manages gRPC channel lifecycle.

    Pools and reuses channels per address. Channels are created lazily
    and closed together on ``close()``.
    """

    def __init__(
        self,
        interceptors: list[grpc.UnaryUnaryClientInterceptor[object, object]] | None = None,
    ) -> None:
        self._channels: dict[str, grpc.Channel] = {}
        self._interceptors = interceptors or []
        self._closed = False

    def get_channel(self, address: str) -> grpc.Channel:
        """Get or create a gRPC channel for the given address."""
        if self._closed:
            raise RuntimeError("ChannelPool is closed")
        if address not in self._channels:
            base_channel = grpc.insecure_channel(address, options=_CHANNEL_OPTIONS)
            channel = grpc.intercept_channel(base_channel, *self._interceptors) if self._interceptors else base_channel
            self._channels[address] = channel
        return self._channels[address]

    def close(self) -> None:
        """Close all pooled channels."""
        if self._closed:
            return
        self._closed = True
        for channel in self._channels.values():
            with contextlib.suppress(Exception):
                channel.close()
        self._channels.clear()

    @property
    def is_closed(self) -> bool:
        return self._closed

    def __enter__(self) -> ChannelPool:
        return self

    def __exit__(self, *args: object) -> None:
        self.close()
