"""Asynchronous KyroDB gRPC client."""

from __future__ import annotations

import asyncio
import logging
from collections.abc import (
    AsyncIterable,
    AsyncIterator,
    Awaitable,
    Callable,
    Iterable,
    Mapping,
    Sequence,
)
from functools import partial
from threading import RLock
from typing import TypeVar

import grpc

from ._converters import (
    to_batch_delete_result,
    to_bulk_load_result,
    to_bulk_query_result,
    to_config_result,
    to_delete_result,
    to_flush_result,
    to_health_result,
    to_insert_ack,
    to_metrics_result,
    to_query_result,
    to_search_response,
    to_snapshot_result,
    to_update_metadata_result,
)
from ._generated import kyrodb_pb2 as pb2
from ._generated import kyrodb_pb2_grpc as pb2_grpc
from .auth import (
    Metadata,
    api_key_metadata,
    merge_metadata,
    validate_api_key,
)
from .client import (
    _USE_CLIENT_DEFAULT_TIMEOUT,
    ApiKeyProvider,
    TimeoutArg,
    TLSConfig,
    _chunked_ids,
    _freeze_call_metadata,
    _is_loopback,
    _merge_channel_options,
    _UseClientDefaultTimeout,
    _validate_doc_id,
    _validate_embedding,
    _validate_tls_config,
)
from .errors import CircuitOpenError, DeadlineExceededError, map_rpc_error
from .models import (
    BatchDeleteResult,
    BulkLoadResult,
    BulkQueryResult,
    ConfigResult,
    DeleteResult,
    FlushResult,
    HealthResult,
    InsertAck,
    InsertItem,
    MetadataFilter,
    MetricsResult,
    QueryResult,
    SearchQuery,
    SearchResponse,
    SnapshotResult,
    UpdateMetadataResult,
)
from .retry import CircuitBreaker, CircuitBreakerPolicy, RetryPolicy, run_with_retry_async

T = TypeVar("T")
AsyncApiKeyProvider = Callable[[], Awaitable[str | None]]
logger = logging.getLogger("kyrodb")


async def _to_async_iterable(items: Iterable[T] | AsyncIterable[T]) -> AsyncIterator[T]:
    if isinstance(items, AsyncIterable):
        async for item in items:
            yield item
    else:
        for item in items:
            yield item


class AsyncKyroDBClient:
    """Async client with full parity for KyroDB v1 gRPC API."""

    def __init__(
        self,
        target: str = "127.0.0.1:50051",
        *,
        api_key: str | None = None,
        tls: TLSConfig | None = None,
        default_namespace: str = "",
        default_timeout_s: float | None = 30.0,
        retry_policy: RetryPolicy | None = None,
        circuit_breaker_policy: CircuitBreakerPolicy | None = None,
        channel_options: Sequence[tuple[str, int | str]] | None = None,
        lb_policy_name: str | None = None,
        compression: grpc.Compression | None = None,
        max_unary_batch_size: int = 10_000,
    ) -> None:
        if default_timeout_s is not None and default_timeout_s <= 0:
            raise ValueError("default_timeout_s must be > 0 when set")
        if max_unary_batch_size <= 0:
            raise ValueError("max_unary_batch_size must be > 0")
        self._target = target
        self._default_timeout_s = default_timeout_s
        self._default_namespace = default_namespace
        self._retry_policy = retry_policy or RetryPolicy()
        self._circuit_breaker = (
            CircuitBreaker(circuit_breaker_policy) if circuit_breaker_policy is not None else None
        )
        self._max_unary_batch_size = max_unary_batch_size
        self._auth_lock = RLock()
        self._dim_lock = RLock()
        self._insert_embedding_dim: int | None = None
        self._api_key: str | None = None
        self._api_key_provider: ApiKeyProvider | None = None
        self._api_key_provider_async: AsyncApiKeyProvider | None = None
        self.set_api_key(api_key)

        if lb_policy_name is not None and not lb_policy_name.strip():
            raise ValueError("lb_policy_name must be non-empty when set")
        options = _merge_channel_options(channel_options)
        if lb_policy_name is not None:
            options.append(("grpc.lb_policy_name", lb_policy_name))

        if tls is None:
            if not _is_loopback(target):
                raise ValueError("TLS is required for non-loopback targets")
            self._channel = grpc.aio.insecure_channel(
                target,
                options=options,
                compression=compression,
            )
        else:
            _validate_tls_config(tls)
            credentials = grpc.ssl_channel_credentials(
                root_certificates=tls.root_certificates,
                private_key=tls.private_key,
                certificate_chain=tls.certificate_chain,
            )
            self._channel = grpc.aio.secure_channel(
                target,
                credentials,
                options=options,
                compression=compression,
            )

        self._stub = pb2_grpc.KyroDBServiceStub(self._channel)
        logger.debug(
            "Initialized AsyncKyroDBClient target=%s tls=%s default_timeout_s=%s",
            target,
            tls is not None,
            default_timeout_s,
        )

    async def close(self) -> None:
        """Close the underlying async gRPC channel."""
        await self._channel.close()
        logger.debug("Closed AsyncKyroDBClient target=%s", self._target)

    async def __aenter__(self) -> AsyncKyroDBClient:
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:  # type: ignore[no-untyped-def]
        await self.close()

    def set_api_key(self, api_key: str | None) -> None:
        """Set or clear static API key used for subsequent calls."""
        validate_api_key(api_key)
        with self._auth_lock:
            self._api_key_provider = None
            self._api_key_provider_async = None
            self._api_key = api_key

    def set_api_key_provider(self, provider: ApiKeyProvider) -> None:
        """Set dynamic API key provider for per-call key rotation."""
        if not callable(provider):
            raise TypeError("provider must be callable")
        probe_key = provider()
        validate_api_key(probe_key)
        with self._auth_lock:
            self._api_key_provider = provider
            self._api_key_provider_async = None
            self._api_key = probe_key

    async def set_api_key_provider_async(self, provider: AsyncApiKeyProvider) -> None:
        """Set async dynamic API key provider for per-call key rotation."""
        if not callable(provider):
            raise TypeError("provider must be callable")
        probe_key = await provider()
        validate_api_key(probe_key)
        with self._auth_lock:
            self._api_key_provider = None
            self._api_key_provider_async = provider
            self._api_key = probe_key

    async def _current_api_key(self) -> str | None:
        with self._auth_lock:
            provider = self._api_key_provider
            provider_async = self._api_key_provider_async
            static_key = self._api_key
        if provider_async is not None:
            current_async = await provider_async()
            validate_api_key(current_async)
            return current_async
        if provider is None:
            return static_key
        current = provider()
        validate_api_key(current)
        return current

    def _record_insert_embedding_dim(self, dim: int) -> None:
        with self._dim_lock:
            previous = self._insert_embedding_dim
            self._insert_embedding_dim = dim
        if previous is not None and previous != dim:
            logger.debug(
                "Insert embedding dimension changed target=%s previous=%d current=%d",
                self._target,
                previous,
                dim,
            )

    def _log_search_embedding_dim_if_mismatch(self, dim: int) -> None:
        with self._dim_lock:
            insert_dim = self._insert_embedding_dim
        if insert_dim is not None and insert_dim != dim:
            logger.debug(
                "Search embedding dimension differs from last insert "
                "target=%s insert_dim=%d search_dim=%d",
                self._target,
                insert_dim,
                dim,
            )

    async def wait_for_ready(
        self,
        *,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
    ) -> None:
        """Wait until channel is ready or deadline is exceeded."""
        timeout = self._timeout(timeout_s)
        try:
            if timeout is None:
                await self._channel.channel_ready()
            else:
                await asyncio.wait_for(self._channel.channel_ready(), timeout=timeout)
            logger.debug("Channel ready target=%s timeout_s=%s", self._target, timeout)
        except (asyncio.TimeoutError, TimeoutError) as exc:
            logger.warning(
                "Channel readiness timed out target=%s timeout_s=%s",
                self._target,
                timeout,
            )
            raise DeadlineExceededError(
                operation="WaitForReady",
                code=grpc.StatusCode.DEADLINE_EXCEEDED,
                details="channel did not become ready before timeout",
                target=self._target,
            ) from exc

    def _timeout(self, timeout_s: TimeoutArg) -> float | None:
        if isinstance(timeout_s, _UseClientDefaultTimeout):
            return self._default_timeout_s
        if timeout_s is not None and timeout_s <= 0:
            raise ValueError("timeout_s must be > 0 when set, or None for unbounded")
        return timeout_s

    async def _metadata(self, metadata: Iterable[tuple[str, str]] | None) -> Metadata:
        return merge_metadata(api_key_metadata(await self._current_api_key()), metadata)

    def _ensure_circuit_allows_call(self, operation: str) -> None:
        if self._circuit_breaker is None:
            return
        allowed, retry_after_s = self._circuit_breaker.allow_call()
        if not allowed:
            raise CircuitOpenError(
                operation=operation,
                code=grpc.StatusCode.UNAVAILABLE,
                details=(
                    "circuit breaker is open"
                    if retry_after_s is None
                    else f"circuit breaker is open; retry after {retry_after_s:.2f}s"
                ),
                target=self._target,
            )

    async def _request_metadata(
        self,
        metadata: Iterable[tuple[str, str]] | None,
        *,
        operation: str,
    ) -> Metadata:
        self._ensure_circuit_allows_call(operation)
        return await self._metadata(_freeze_call_metadata(metadata))

    async def _safe_call(
        self,
        fn: Callable[[], Awaitable[T]],
        operation: str,
        retry: bool,
    ) -> T:
        self._ensure_circuit_allows_call(operation)
        try:
            if retry:
                response = await run_with_retry_async(fn, self._retry_policy)
            else:
                response = await fn()
            if self._circuit_breaker is not None:
                self._circuit_breaker.record_success()
            return response
        except grpc.RpcError as exc:
            code = exc.code() if hasattr(exc, "code") else grpc.StatusCode.UNKNOWN
            if self._circuit_breaker is not None:
                self._circuit_breaker.record_failure(code)
            safe_auth_metadata = (("x-api-key", "<redacted>"),)
            logger.warning(
                "RPC failed operation=%s target=%s code=%s retry=%s metadata=%s",
                operation,
                self._target,
                code.name if hasattr(code, "name") else str(code),
                retry,
                safe_auth_metadata,
            )
            raise map_rpc_error(exc, operation=operation, target=self._target) from exc

    async def insert(
        self,
        *,
        doc_id: int,
        embedding: Sequence[float],
        metadata: Mapping[str, str] | None = None,
        namespace: str | None = None,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> InsertAck:
        """Insert a single document vector into KyroDB."""
        _validate_doc_id(doc_id)
        vector = _validate_embedding(embedding, "embedding")
        request = pb2.InsertRequest(
            doc_id=doc_id,
            embedding=vector,
            metadata=dict(metadata or {}),
            namespace=namespace if namespace is not None else self._default_namespace,
        )
        self._record_insert_embedding_dim(len(vector))
        request_metadata = await self._request_metadata(call_metadata, operation="Insert")
        response = await self._safe_call(
            lambda: self._stub.Insert(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="Insert",
            retry=False,
        )
        return to_insert_ack(response)

    async def bulk_insert(
        self,
        items: Iterable[InsertItem] | AsyncIterable[InsertItem],
        *,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> InsertAck:
        """Stream-insert multiple documents."""
        request_metadata = await self._request_metadata(call_metadata, operation="BulkInsert")

        async def req_stream() -> AsyncIterator[pb2.InsertRequest]:
            async for item in _to_async_iterable(items):
                _validate_doc_id(item.doc_id)
                vector = _validate_embedding(item.embedding, "embedding")
                self._record_insert_embedding_dim(len(vector))
                yield pb2.InsertRequest(
                    doc_id=item.doc_id,
                    embedding=vector,
                    metadata=dict(item.metadata),
                    namespace=item.namespace or self._default_namespace,
                )

        response = await self._safe_call(
            lambda: self._stub.BulkInsert(
                req_stream(),
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="BulkInsert",
            retry=False,
        )
        return to_insert_ack(response)

    async def bulk_load_hnsw(
        self,
        items: Iterable[InsertItem] | AsyncIterable[InsertItem],
        *,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> BulkLoadResult:
        """Bulk-load vectors directly into HNSW."""
        request_metadata = await self._request_metadata(call_metadata, operation="BulkLoadHnsw")

        async def req_stream() -> AsyncIterator[pb2.InsertRequest]:
            async for item in _to_async_iterable(items):
                _validate_doc_id(item.doc_id)
                vector = _validate_embedding(item.embedding, "embedding")
                self._record_insert_embedding_dim(len(vector))
                yield pb2.InsertRequest(
                    doc_id=item.doc_id,
                    embedding=vector,
                    metadata=dict(item.metadata),
                    namespace=item.namespace or self._default_namespace,
                )

        response = await self._safe_call(
            lambda: self._stub.BulkLoadHnsw(
                req_stream(),
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="BulkLoadHnsw",
            retry=False,
        )
        return to_bulk_load_result(response)

    async def delete(
        self,
        *,
        doc_id: int,
        namespace: str | None = None,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> DeleteResult:
        """Delete a document by ID."""
        _validate_doc_id(doc_id)
        request = pb2.DeleteRequest(
            doc_id=doc_id,
            namespace=namespace if namespace is not None else self._default_namespace,
        )
        request_metadata = await self._request_metadata(call_metadata, operation="Delete")
        response = await self._safe_call(
            lambda: self._stub.Delete(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="Delete",
            retry=False,
        )
        return to_delete_result(response)

    async def update_metadata(
        self,
        *,
        doc_id: int,
        metadata: Mapping[str, str],
        merge: bool = True,
        namespace: str | None = None,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> UpdateMetadataResult:
        """Update metadata for an existing document."""
        _validate_doc_id(doc_id)
        request = pb2.UpdateMetadataRequest(
            doc_id=doc_id,
            metadata=dict(metadata),
            merge=merge,
            namespace=namespace if namespace is not None else self._default_namespace,
        )
        request_metadata = await self._request_metadata(call_metadata, operation="UpdateMetadata")
        response = await self._safe_call(
            lambda: self._stub.UpdateMetadata(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="UpdateMetadata",
            retry=False,
        )
        return to_update_metadata_result(response)

    async def query(
        self,
        *,
        doc_id: int,
        include_embedding: bool = False,
        namespace: str | None = None,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> QueryResult:
        """Fetch a document by ID."""
        _validate_doc_id(doc_id)
        request = pb2.QueryRequest(
            doc_id=doc_id,
            include_embedding=include_embedding,
            namespace=namespace if namespace is not None else self._default_namespace,
        )
        request_metadata = await self._request_metadata(call_metadata, operation="Query")
        response = await self._safe_call(
            lambda: self._stub.Query(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="Query",
            retry=True,
        )
        return to_query_result(response)

    async def search(
        self,
        *,
        query_embedding: Sequence[float],
        k: int,
        min_score: float = 0.0,
        namespace: str | None = None,
        include_embeddings: bool = False,
        filter: MetadataFilter | None = None,
        ef_search: int = 0,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> SearchResponse:
        """Run k-NN search against KyroDB."""
        if k <= 0 or k > 1000:
            raise ValueError("k must be between 1 and 1000")
        if ef_search < 0:
            raise ValueError("ef_search must be >= 0")
        vector = _validate_embedding(query_embedding, "query_embedding")
        self._log_search_embedding_dim_if_mismatch(len(vector))
        request = pb2.SearchRequest(
            query_embedding=vector,
            k=k,
            min_score=min_score,
            namespace=namespace if namespace is not None else self._default_namespace,
            include_embeddings=include_embeddings,
            ef_search=ef_search,
        )
        if filter is not None:
            request.filter.CopyFrom(filter.to_proto())
        request_metadata = await self._request_metadata(call_metadata, operation="Search")

        response = await self._safe_call(
            lambda: self._stub.Search(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="Search",
            retry=True,
        )
        return to_search_response(response)

    async def bulk_search(
        self,
        requests: Iterable[SearchQuery] | AsyncIterable[SearchQuery],
        *,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> AsyncIterator[SearchResponse]:
        """Run streaming batch search and yield responses."""
        request_metadata = await self._request_metadata(call_metadata, operation="BulkSearch")

        async def req_stream() -> AsyncIterator[pb2.SearchRequest]:
            async for query in _to_async_iterable(requests):
                if query.k <= 0 or query.k > 1000:
                    raise ValueError("k must be between 1 and 1000")
                if query.ef_search < 0:
                    raise ValueError("ef_search must be >= 0")
                vector = _validate_embedding(query.query_embedding, "query_embedding")
                self._log_search_embedding_dim_if_mismatch(len(vector))
                request = pb2.SearchRequest(
                    query_embedding=vector,
                    k=query.k,
                    min_score=query.min_score,
                    namespace=query.namespace or self._default_namespace,
                    include_embeddings=query.include_embeddings,
                    ef_search=query.ef_search,
                )
                if query.filter is not None:
                    request.filter.CopyFrom(query.filter.to_proto())
                yield request

        try:
            response_stream = self._stub.BulkSearch(
                req_stream(),
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            )
            async for response in response_stream:
                yield to_search_response(response)
            if self._circuit_breaker is not None:
                self._circuit_breaker.record_success()
        except grpc.RpcError as exc:
            code = exc.code() if hasattr(exc, "code") else grpc.StatusCode.UNKNOWN
            if self._circuit_breaker is not None:
                self._circuit_breaker.record_failure(code)
            raise map_rpc_error(exc, operation="BulkSearch", target=self._target) from exc

    async def bulk_query(
        self,
        *,
        doc_ids: Sequence[int],
        include_embeddings: bool = False,
        namespace: str | None = None,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> BulkQueryResult:
        """Fetch multiple documents by ID in one request."""
        if not doc_ids:
            raise ValueError("doc_ids must be non-empty")
        for doc_id in doc_ids:
            _validate_doc_id(doc_id)
        request_metadata = await self._request_metadata(call_metadata, operation="BulkQuery")
        if len(doc_ids) <= self._max_unary_batch_size:
            request = pb2.BulkQueryRequest(
                doc_ids=list(doc_ids),
                include_embeddings=include_embeddings,
                namespace=namespace if namespace is not None else self._default_namespace,
            )
            response = await self._safe_call(
                lambda: self._stub.BulkQuery(
                    request,
                    timeout=self._timeout(timeout_s),
                    metadata=request_metadata,
                ),
                operation="BulkQuery",
                retry=True,
            )
            return to_bulk_query_result(response)

        results: list[QueryResult] = []
        total_found = 0
        total_requested = 0
        errors: list[str] = []
        for chunk in _chunked_ids(doc_ids, self._max_unary_batch_size):
            request = pb2.BulkQueryRequest(
                doc_ids=chunk,
                include_embeddings=include_embeddings,
                namespace=namespace if namespace is not None else self._default_namespace,
            )
            call = partial(
                self._stub.BulkQuery,
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            )
            response = await self._safe_call(
                call,
                operation="BulkQuery",
                retry=True,
            )
            converted = to_bulk_query_result(response)
            results.extend(converted.results)
            total_found += converted.total_found
            total_requested += converted.total_requested
            if converted.error:
                errors.append(converted.error)

        return BulkQueryResult(
            results=tuple(results),
            total_found=total_found,
            total_requested=total_requested,
            error=" | ".join(errors),
        )

    async def batch_delete_ids(
        self,
        *,
        doc_ids: Sequence[int],
        namespace: str | None = None,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> BatchDeleteResult:
        """Delete multiple documents by explicit IDs."""
        if not doc_ids:
            raise ValueError("doc_ids must be non-empty")
        for doc_id in doc_ids:
            _validate_doc_id(doc_id)
        request_metadata = await self._request_metadata(call_metadata, operation="BatchDelete")

        if len(doc_ids) <= self._max_unary_batch_size:
            request = pb2.BatchDeleteRequest(
                ids=pb2.IdList(doc_ids=list(doc_ids)),
                namespace=namespace if namespace is not None else self._default_namespace,
            )
            response = await self._safe_call(
                lambda: self._stub.BatchDelete(
                    request,
                    timeout=self._timeout(timeout_s),
                    metadata=request_metadata,
                ),
                operation="BatchDelete",
                retry=False,
            )
            return to_batch_delete_result(response)

        deleted_total = 0
        errors: list[str] = []
        all_success = True
        for chunk in _chunked_ids(doc_ids, self._max_unary_batch_size):
            request = pb2.BatchDeleteRequest(
                ids=pb2.IdList(doc_ids=chunk),
                namespace=namespace if namespace is not None else self._default_namespace,
            )
            call = partial(
                self._stub.BatchDelete,
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            )
            response = await self._safe_call(
                call,
                operation="BatchDelete",
                retry=False,
            )
            converted = to_batch_delete_result(response)
            deleted_total += converted.deleted_count
            all_success = all_success and converted.success
            if converted.error:
                errors.append(converted.error)

        return BatchDeleteResult(
            success=all_success,
            deleted_count=deleted_total,
            error=" | ".join(errors),
        )

    async def health(
        self,
        *,
        component: str = "",
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> HealthResult:
        """Fetch service health status."""
        request = pb2.HealthRequest(component=component)
        request_metadata = await self._request_metadata(call_metadata, operation="Health")
        response = await self._safe_call(
            lambda: self._stub.Health(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="Health",
            retry=True,
        )
        return to_health_result(response)

    async def metrics(
        self,
        *,
        categories: Sequence[str] | None = None,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> MetricsResult:
        """Fetch service metrics snapshot."""
        request = pb2.MetricsRequest(categories=list(categories or ()))
        request_metadata = await self._request_metadata(call_metadata, operation="Metrics")
        response = await self._safe_call(
            lambda: self._stub.Metrics(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="Metrics",
            retry=True,
        )
        return to_metrics_result(response)

    async def flush_hot_tier(
        self,
        *,
        force: bool = True,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> FlushResult:
        """Force flush hot tier to cold tier."""
        request = pb2.FlushRequest(force=force)
        request_metadata = await self._request_metadata(call_metadata, operation="FlushHotTier")
        response = await self._safe_call(
            lambda: self._stub.FlushHotTier(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="FlushHotTier",
            retry=False,
        )
        return to_flush_result(response)

    async def create_snapshot(
        self,
        *,
        path: str = "",
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> SnapshotResult:
        """Create a snapshot on the server."""
        request = pb2.SnapshotRequest(path=path)
        request_metadata = await self._request_metadata(call_metadata, operation="CreateSnapshot")
        response = await self._safe_call(
            lambda: self._stub.CreateSnapshot(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="CreateSnapshot",
            retry=False,
        )
        return to_snapshot_result(response)

    async def get_config(
        self,
        *,
        timeout_s: TimeoutArg = _USE_CLIENT_DEFAULT_TIMEOUT,
        call_metadata: Iterable[tuple[str, str]] | None = None,
    ) -> ConfigResult:
        """Fetch effective server configuration."""
        request = pb2.ConfigRequest()
        request_metadata = await self._request_metadata(call_metadata, operation="GetConfig")
        response = await self._safe_call(
            lambda: self._stub.GetConfig(
                request,
                timeout=self._timeout(timeout_s),
                metadata=request_metadata,
            ),
            operation="GetConfig",
            retry=True,
        )
        return to_config_result(response)
