import asyncio
import base64
import json
import socket
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union
from urllib import parse

import grpc
from grpc._cython import cygrpc

from pymilvus.client.call_context import CallContext, _api_level_md
from pymilvus.client.types import GrantInfo, ResourceGroupConfig
from pymilvus.decorators import (
    ignore_unimplemented,
    retry_on_rpc_failure,
    retry_on_schema_mismatch,
)
from pymilvus.exceptions import (
    AmbiguousIndexName,
    DescribeCollectionException,
    ErrorCode,
    ExceptionsMessage,
    MilvusException,
    ParamError,
)
from pymilvus.grpc_gen import common_pb2, milvus_pb2_grpc
from pymilvus.grpc_gen import milvus_pb2 as milvus_types
from pymilvus.orm.schema import Function, Highlighter
from pymilvus.settings import Config

from . import entity_helper, ts_utils, utils
from .abstract import AnnSearchRequest, BaseRanker, CollectionSchema, FieldSchema, MutationResult
from .async_interceptor import async_header_adder_interceptor
from .cache import GlobalCache
from .check import (
    check_id_and_data,
    check_pass_param,
    is_legal_host,
    is_legal_port,
)
from .constants import ITERATOR_SESSION_TS_FIELD
from .embedding_list import EmbeddingList
from .prepare import Prepare
from .search_result import SearchResult
from .types import (
    AnalyzeResult,
    CompactionState,
    DatabaseInfo,
    HybridExtraList,
    IndexState,
    LoadState,
    ReplicaInfo,
    Shard,
    State,
    Status,
    get_extra_info,
)
from .utils import (
    check_invalid_binary_vector,
    check_status,
    get_server_type,
    is_successful,
    len_of,
)


class AsyncGrpcHandler:
    def __init__(
        self,
        uri: str = Config.GRPC_URI,
        host: str = "",
        port: str = "",
        channel: Optional[grpc.aio.Channel] = None,
        **kwargs,
    ) -> None:
        self._async_stub = None
        self._async_channel = channel

        addr = kwargs.get("address")
        self._address = addr if addr is not None else self.__get_address(uri, host, port)
        self._log_level = None
        self._user = kwargs.get("user")
        self._grpc_options = kwargs.get("grpc_options", {})
        self._set_authorization(**kwargs)
        self._setup_grpc_channel(**kwargs)
        self._is_channel_ready = False
        self.callbacks = []  # Do nothing
        self._server_info_cache = None

    def __get_address(self, uri: str, host: str, port: str) -> str:
        if host != "" and port != "" and is_legal_host(host) and is_legal_port(port):
            return f"{host}:{port}"

        try:
            parsed_uri = parse.urlparse(uri)
        except Exception as e:
            raise ParamError(message=f"Illegal uri: [{uri}], {e}") from e
        return parsed_uri.netloc

    def _set_authorization(self, **kwargs):
        secure = kwargs.get("secure", False)
        if not isinstance(secure, bool):
            raise ParamError(message="secure must be bool type")
        self._secure = secure
        self._client_pem_path = kwargs.get("client_pem_path", "")
        self._client_key_path = kwargs.get("client_key_path", "")
        self._ca_pem_path = kwargs.get("ca_pem_path", "")
        self._server_pem_path = kwargs.get("server_pem_path", "")
        self._server_name = kwargs.get("server_name", "")

        self._async_authorization_interceptor = None

    def __enter__(self):
        return self

    def __exit__(self: object, exc_type: object, exc_val: object, exc_tb: object):
        pass

    async def close(self):
        await self._async_channel.close()
        self._async_channel = None

    def _setup_authorization_interceptor(self, user: str, password: str, token: str):
        keys = []
        values = []
        if token:
            authorization = base64.b64encode(f"{token}".encode())
            keys.append("authorization")
            values.append(authorization)
        elif user and password:
            authorization = base64.b64encode(f"{user}:{password}".encode())
            keys.append("authorization")
            values.append(authorization)
        if len(keys) > 0 and len(values) > 0:
            self._async_authorization_interceptor = async_header_adder_interceptor(keys, values)
            self._final_channel._unary_unary_interceptors.append(
                self._async_authorization_interceptor
            )

    def _setup_grpc_channel(self, **kwargs):
        if self._async_channel is None:
            # Default gRPC options
            default_opts = {
                cygrpc.ChannelArgKey.max_send_message_length: -1,
                cygrpc.ChannelArgKey.max_receive_message_length: -1,
                "grpc.enable_retries": 1,
                "grpc.keepalive_time_ms": 10000,
                "grpc.keepalive_timeout_ms": 5000,
                "grpc.keepalive_permit_without_calls": True,
            }
            # Merge user-provided options (user options override defaults)
            default_opts.update(self._grpc_options)
            opts = list(default_opts.items())
            if not self._secure:
                self._async_channel = grpc.aio.insecure_channel(
                    self._address,
                    options=opts,
                )
            else:
                if self._server_name != "":
                    opts.append(("grpc.ssl_target_name_override", self._server_name))

                root_cert, private_k, cert_chain = None, None, None
                if self._server_pem_path != "":
                    with Path(self._server_pem_path).open("rb") as f:
                        root_cert = f.read()
                elif (
                    self._client_pem_path != ""
                    and self._client_key_path != ""
                    and self._ca_pem_path != ""
                ):
                    with Path(self._ca_pem_path).open("rb") as f:
                        root_cert = f.read()
                    with Path(self._client_key_path).open("rb") as f:
                        private_k = f.read()
                    with Path(self._client_pem_path).open("rb") as f:
                        cert_chain = f.read()

                creds = grpc.ssl_channel_credentials(
                    root_certificates=root_cert,
                    private_key=private_k,
                    certificate_chain=cert_chain,
                )
                self._async_channel = grpc.aio.secure_channel(
                    self._address,
                    creds,
                    options=opts,
                )

        # avoid to add duplicate headers.
        self._final_channel = self._async_channel

        if self._async_authorization_interceptor:
            self._final_channel._unary_unary_interceptors.append(
                self._async_authorization_interceptor
            )
        else:
            self._setup_authorization_interceptor(
                kwargs.get("user"),
                kwargs.get("password"),
                kwargs.get("token"),
            )
        if self._log_level:
            async_log_level_interceptor = async_header_adder_interceptor(
                ["log-level"], [self._log_level]
            )
            self._final_channel._unary_unary_interceptors.append(async_log_level_interceptor)
            self._log_level = None
        self._async_stub = milvus_pb2_grpc.MilvusServiceStub(self._final_channel)

    @property
    def server_address(self):
        return self._address

    def get_server_type(self):
        return get_server_type(self.server_address.split(":")[0])

    async def ensure_channel_ready(self):
        try:
            if not self._is_channel_ready:
                # wait for channel ready
                await self._async_channel.channel_ready()
                # set identifier interceptor
                host = socket.gethostname()
                req = Prepare.register_request(self._user, host)
                response = await self._async_stub.Connect(request=req)
                check_status(response.status)
                _async_identifier_interceptor = async_header_adder_interceptor(
                    ["identifier"], [str(response.identifier)]
                )
                self._async_channel._unary_unary_interceptors.append(_async_identifier_interceptor)
                self._async_stub = milvus_pb2_grpc.MilvusServiceStub(self._async_channel)

                self._is_channel_ready = True
        except grpc.FutureTimeoutError as e:
            raise MilvusException(
                code=Status.CONNECT_FAILED,
                message=f"Fail connecting to server on {self._address}, illegal connection params or server unavailable",
            ) from e
        except Exception as e:
            raise e from e

    @retry_on_rpc_failure()
    async def create_collection(
        self,
        collection_name: str,
        fields: List,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.create_collection_request(collection_name, fields, **kwargs)
        response = await self._async_stub.CreateCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def drop_collection(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.drop_collection_request(collection_name)
        response = await self._async_stub.DropCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)
        # Invalidate global schema cache
        self._invalidate_schema(
            collection_name,
            db_name=(context.get_db_name() if context else kwargs.get("db_name", "")),
        )

    @retry_on_rpc_failure()
    async def truncate_collection(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.truncate_collection_request(collection_name)
        response = await self._async_stub.TruncateCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)

    @retry_on_rpc_failure()
    async def load_collection(
        self,
        collection_name: str,
        replica_number: Optional[int] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()

        check_pass_param(timeout=timeout)
        request = Prepare.load_collection(collection_name, replica_number, **kwargs)
        response = await self._async_stub.LoadCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

        await self.wait_for_loading_collection(
            collection_name=collection_name,
            is_refresh=request.refresh,
            timeout=timeout,
            context=context,
            **kwargs,
        )

    @retry_on_rpc_failure()
    async def wait_for_loading_collection(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        is_refresh: bool = False,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        start = time.time()

        def can_loop(t: int) -> bool:
            return True if timeout is None else t <= (start + timeout)

        while can_loop(time.time()):
            progress = await self.get_loading_progress(
                collection_name=collection_name,
                is_refresh=is_refresh,
                timeout=timeout,
                context=context,
                **kwargs,
            )
            if progress >= 100:
                return
            await asyncio.sleep(Config.WaitTimeDurationWhenLoad)
        raise MilvusException(
            message=f"wait for loading collection timeout, collection: {collection_name}"
        )

    @retry_on_rpc_failure()
    async def get_loading_progress(
        self,
        collection_name: str,
        partition_names: Optional[List[str]] = None,
        timeout: Optional[float] = None,
        is_refresh: bool = False,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        request = Prepare.get_loading_progress(collection_name, partition_names)
        response = await self._async_stub.GetLoadingProgress(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        if is_refresh:
            return response.refresh_progress
        return response.progress

    @retry_on_rpc_failure()
    async def describe_collection(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.describe_collection_request(collection_name)
        response = await self._async_stub.DescribeCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        status = response.status

        if is_successful(status):
            return CollectionSchema(raw=response).dict()

        raise DescribeCollectionException(status.code, status.reason, status.error_code)

    @retry_on_rpc_failure()
    async def has_collection(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> bool:
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.describe_collection_request(collection_name)
        reply = await self._async_stub.DescribeCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )

        if (
            reply.status.error_code == common_pb2.UnexpectedError
            and "can't find collection" in reply.status.reason
        ):
            return False

        if reply.status.error_code == common_pb2.CollectionNotExists:
            return False

        if is_successful(reply.status):
            return True

        if reply.status.code == ErrorCode.COLLECTION_NOT_FOUND:
            return False

        raise MilvusException(reply.status.code, reply.status.reason, reply.status.error_code)

    @retry_on_rpc_failure()
    async def list_collections(
        self, timeout: Optional[float] = None, context: Optional[CallContext] = None, **kwargs
    ) -> List[str]:
        await self.ensure_channel_ready()
        request = Prepare.show_collections_request()
        response = await self._async_stub.ShowCollections(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        status = response.status
        check_status(status)
        return list(response.collection_names)

    @retry_on_rpc_failure()
    async def get_collection_stats(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        index_param = Prepare.get_collection_stats_request(collection_name)
        response = await self._async_stub.GetCollectionStatistics(
            index_param, timeout=timeout, metadata=_api_level_md(context)
        )
        status = response.status
        check_status(status)
        return response.stats

    @retry_on_rpc_failure()
    async def get_partition_stats(
        self,
        collection_name: str,
        partition_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        req = Prepare.get_partition_stats_request(collection_name, partition_name)
        response = await self._async_stub.GetPartitionStatistics(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        status = response.status
        check_status(status)
        return response.stats

    @retry_on_rpc_failure()
    async def get_load_state(
        self,
        collection_name: str,
        partition_names: Optional[List[str]] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        request = Prepare.get_load_state(collection_name, partition_names)
        response = await self._async_stub.GetLoadState(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return LoadState(response.state)

    @retry_on_rpc_failure()
    async def refresh_load(
        self,
        collection_name: str,
        partition_names: Optional[List[str]] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        request = Prepare.get_loading_progress(collection_name, partition_names)
        response = await self._async_stub.GetLoadingProgress(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return response.refresh_progress

    @retry_on_rpc_failure()
    async def get_server_version(
        self,
        timeout: Optional[float] = None,
        detail: bool = False,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> Union[str, dict]:
        await self.ensure_channel_ready()
        if detail:
            if self._server_info_cache is None:
                req = Prepare.register_request("", "")
                resp = await self._async_stub.Connect(request=req, timeout=timeout)
                check_status(resp.status)
                info = resp.server_info
                self._server_info_cache = {
                    "version": info.build_tags,
                    "build_time": info.build_time,
                    "git_commit": info.git_commit,
                    "go_version": info.go_version,
                    "deploy_mode": info.deploy_mode,
                }
            return self._server_info_cache
        req = Prepare.get_server_version()
        resp = await self._async_stub.GetVersion(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return resp.version

    @retry_on_rpc_failure()
    async def describe_replica(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> List[ReplicaInfo]:
        await self.ensure_channel_ready()
        schema, _ = await self._get_schema(
            collection_name, timeout=timeout, context=context, **kwargs
        )
        collection_id = schema["collection_id"]

        req = Prepare.get_replicas(collection_id)
        response = await self._async_stub.GetReplicas(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)

        groups = []
        for replica in response.replicas:
            shards = [
                Shard(s.dm_channel_name, s.node_ids, s.leaderID) for s in replica.shard_replicas
            ]
            groups.append(
                ReplicaInfo(
                    replica.replicaID,
                    shards,
                    replica.node_ids,
                    replica.resource_group_name,
                    replica.num_outbound_node,
                )
            )

        return groups

    @retry_on_rpc_failure()
    async def rename_collection(
        self,
        old_name: str,
        new_name: str,
        new_db_name: str = "",
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=new_name, timeout=timeout)
        check_pass_param(collection_name=old_name)
        if new_db_name:
            check_pass_param(db_name=new_db_name)
        request = Prepare.rename_collections_request(old_name, new_name, new_db_name)
        status = await self._async_stub.RenameCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    async def _get_info(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        schema = kwargs.get("schema")
        schema, schema_timestamp = await self._get_schema(
            collection_name, timeout=timeout, context=context, **kwargs
        )
        fields_info = schema.get("fields")
        struct_fields_info = schema.get("struct_array_fields", [])
        enable_dynamic = schema.get("enable_dynamic_field", False)

        return fields_info, struct_fields_info, enable_dynamic, schema_timestamp

    async def _get_schema(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> Tuple[dict, int]:
        """
        Get collection schema, using cache when available.

        Returns:
            Tuple of (schema_dict, schema_timestamp)
        """
        cache = GlobalCache.schema
        endpoint = self.server_address
        db_name = context.get_db_name() if context else ""

        cached = cache.get(endpoint, db_name, collection_name)
        if cached is not None:
            return cached, cached.get("update_timestamp", 0)

        # Fetch from server and cache
        schema = await self.describe_collection(
            collection_name, timeout=timeout, context=context, **kwargs
        )
        cache.set(endpoint, db_name, collection_name, schema)
        return schema, schema.get("update_timestamp", 0)

    def _invalidate_schema(self, collection_name: str, db_name: str = "") -> None:
        """Invalidate cached schema for a collection."""
        GlobalCache.schema.invalidate(self.server_address, db_name, collection_name)

    def _invalidate_db_schemas(self, db_name: str) -> None:
        """Invalidate all cached schemas for a database."""
        GlobalCache.schema.invalidate_db(self.server_address, db_name)

    @retry_on_rpc_failure()
    async def release_collection(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.release_collection("", collection_name)
        response = await self._async_stub.ReleaseCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    @retry_on_schema_mismatch()
    async def insert_rows(
        self,
        collection_name: str,
        entities: Union[Dict, List[Dict]],
        partition_name: Optional[str] = None,
        schema: Optional[dict] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        request = await self._prepare_row_insert_request(
            collection_name, entities, partition_name, schema, timeout, context=context, **kwargs
        )
        resp = await self._async_stub.Insert(
            request=request,
            timeout=timeout,
            metadata=_api_level_md(context),
        )
        check_status(resp.status)
        ts_utils.update_collection_ts(
            collection_name,
            resp.timestamp,
            self.server_address,
            (context.get_db_name() if context else ""),
        )
        return MutationResult(resp)

    async def _prepare_row_insert_request(
        self,
        collection_name: str,
        entity_rows: Union[List[Dict], Dict],
        partition_name: Optional[str] = None,
        schema: Optional[dict] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        if isinstance(entity_rows, dict):
            entity_rows = [entity_rows]

        schema, schema_timestamp = await self._get_schema(
            collection_name, schema=schema, timeout=timeout, context=context, **kwargs
        )
        fields_info = schema.get("fields")
        struct_fields_info = schema.get("struct_array_fields", [])  # Default to empty list
        enable_dynamic = schema.get("enable_dynamic_field", False)
        namespace = kwargs.get("namespace")

        return Prepare.row_insert_param(
            collection_name,
            entity_rows,
            partition_name,
            fields_info,
            struct_fields_info,
            enable_dynamic=enable_dynamic,
            schema_timestamp=schema_timestamp,
            namespace=namespace,
        )

    @retry_on_rpc_failure()
    async def get_persistent_segment_infos(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> List[milvus_types.PersistentSegmentInfo]:
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        req = Prepare.get_persistent_segment_info_request(collection_name)
        response = await self._async_stub.GetPersistentSegmentInfo(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return response.infos

    async def delete(
        self,
        collection_name: str,
        expression: str,
        partition_name: Optional[str] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        try:
            req = Prepare.delete_request(
                collection_name=collection_name,
                filter=expression,
                partition_name=partition_name,
                consistency_level=kwargs.pop("consistency_level", 0),
                **kwargs,
            )

            response = await self._async_stub.Delete(
                req, timeout=timeout, metadata=_api_level_md(context)
            )

            m = MutationResult(response)
            ts_utils.update_collection_ts(
                collection_name,
                m.timestamp,
                self.server_address,
                (context.get_db_name() if context else ""),
            )
        except Exception as err:
            raise err from err
        else:
            return m

    async def _prepare_batch_upsert_request(
        self,
        collection_name: str,
        entities: List,
        partition_name: Optional[str] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        param = kwargs.get("upsert_param")
        if param and not isinstance(param, milvus_types.UpsertRequest):
            raise ParamError(message="The value of key 'upsert_param' is invalid")
        if not isinstance(entities, list):
            raise ParamError(message="'entities' must be a list, please provide valid entity data.")

        # Extract partial_update parameter from kwargs
        partial_update = kwargs.get("partial_update", False)

        schema = kwargs.get("schema")
        schema, _ = await self._get_schema(
            collection_name, schema=schema, timeout=timeout, context=context, **kwargs
        )

        fields_info = schema["fields"]

        return (
            param
            if param
            else Prepare.batch_upsert_param(
                collection_name,
                entities,
                partition_name,
                fields_info,
                partial_update=partial_update,
            )
        )

    @retry_on_rpc_failure()
    async def upsert(
        self,
        collection_name: str,
        entities: List,
        partition_name: Optional[str] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        if not check_invalid_binary_vector(entities):
            raise ParamError(message="Invalid binary vector data exists")

        request = await self._prepare_batch_upsert_request(
            collection_name, entities, partition_name, timeout, context=context, **kwargs
        )
        response = await self._async_stub.Upsert(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        m = MutationResult(response)
        ts_utils.update_collection_ts(
            collection_name,
            m.timestamp,
            self.server_address,
            (context.get_db_name() if context else ""),
        )
        return m

    async def _prepare_row_upsert_request(
        self,
        collection_name: str,
        rows: List,
        partition_name: Optional[str] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        if not isinstance(rows, list):
            raise ParamError(message="'rows' must be a list, please provide valid row data.")

        partial_update = kwargs.get("partial_update", False)
        fields_info, struct_fields_info, enable_dynamic, schema_timestamp = await self._get_info(
            collection_name, timeout, context, **kwargs
        )
        return Prepare.row_upsert_param(
            collection_name,
            rows,
            partition_name,
            fields_info,
            struct_fields_info,
            enable_dynamic=enable_dynamic,
            partial_update=partial_update,
            schema_timestamp=schema_timestamp,
        )

    @retry_on_rpc_failure()
    @retry_on_schema_mismatch()
    async def upsert_rows(
        self,
        collection_name: str,
        entities: List,
        partition_name: Optional[str] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        if isinstance(entities, dict):
            entities = [entities]
        request = await self._prepare_row_upsert_request(
            collection_name, entities, partition_name, timeout, context=context, **kwargs
        )
        response = await self._async_stub.Upsert(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        m = MutationResult(response)
        ts_utils.update_collection_ts(
            collection_name,
            m.timestamp,
            self.server_address,
            (context.get_db_name() if context else ""),
        )
        return m

    async def _execute_search(
        self,
        request: milvus_types.SearchRequest,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        response = await self._async_stub.Search(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        round_decimal = kwargs.get("round_decimal", -1)
        return SearchResult(
            response.results,
            round_decimal,
            status=response.status,
            session_ts=response.session_ts,
        )

    async def _execute_hybrid_search(
        self,
        request: milvus_types.HybridSearchRequest,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        response = await self._async_stub.HybridSearch(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        round_decimal = kwargs.get("round_decimal", -1)
        return SearchResult(response.results, round_decimal, status=response.status)

    @retry_on_rpc_failure()
    async def search(
        self,
        collection_name: str,
        anns_field: str,
        param: Dict,
        limit: int,
        data: Optional[Union[List[List[float]], utils.SparseMatrixInputType]] = None,
        ids: Optional[Union[List[int], List[str], str, int]] = None,
        expression: Optional[str] = None,
        partition_names: Optional[List[str]] = None,
        output_fields: Optional[List[str]] = None,
        round_decimal: int = -1,
        timeout: Optional[float] = None,
        ranker: Optional[Function] = None,
        highlighter: Optional[Highlighter] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        if isinstance(ids, (int, str)):
            ids = [ids]
        check_id_and_data(ids, data)

        check_pass_param(
            limit=limit,
            round_decimal=round_decimal,
            anns_field=anns_field,
            partition_name_array=partition_names,
            output_fields=output_fields,
            guarantee_timestamp=kwargs.get("guarantee_timestamp"),
            timeout=timeout,
        )

        # Convert EmbeddingList to flat array if present
        if isinstance(data, list) and len(data) > 0 and isinstance(data[0], EmbeddingList):
            data = [emb_list.to_flat_array() for emb_list in data]
            kwargs["is_embedding_list"] = True

        use_default_consistency = ts_utils.construct_guarantee_ts(
            collection_name, kwargs, self.server_address, (context.get_db_name() if context else "")
        )

        request = Prepare.search_requests_with_expr(
            collection_name=collection_name,
            data=data,
            ids=ids,
            anns_field=anns_field,
            param=param,
            limit=limit,
            expr=expression,
            partition_names=partition_names,
            output_fields=output_fields,
            round_decimal=round_decimal,
            ranker=ranker,
            highlighter=highlighter,
            use_default_consistency=use_default_consistency,
            **kwargs,
        )
        return await self._execute_search(
            request, timeout, context=context, round_decimal=round_decimal, **kwargs
        )

    @retry_on_rpc_failure()
    async def hybrid_search(
        self,
        collection_name: str,
        reqs: List[AnnSearchRequest],
        rerank: Union[BaseRanker, Function],
        limit: int,
        partition_names: Optional[List[str]] = None,
        output_fields: Optional[List[str]] = None,
        round_decimal: int = -1,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(
            limit=limit,
            round_decimal=round_decimal,
            partition_name_array=partition_names,
            output_fields=output_fields,
            guarantee_timestamp=kwargs.get("guarantee_timestamp"),
            timeout=timeout,
        )

        use_default_consistency = ts_utils.construct_guarantee_ts(
            collection_name, kwargs, self.server_address, (context.get_db_name() if context else "")
        )

        requests = []
        for req in reqs:
            data = req.data
            req_kwargs = dict(kwargs)
            # Convert EmbeddingList to flat array if present
            if isinstance(data, list) and len(data) > 0 and isinstance(data[0], EmbeddingList):
                data = [emb_list.to_flat_array() for emb_list in data]
                req_kwargs["is_embedding_list"] = True

            search_request = Prepare.search_requests_with_expr(
                collection_name=collection_name,
                data=data,
                anns_field=req.anns_field,
                param=req.param,
                limit=req.limit,
                expr=req.expr,
                partition_names=partition_names,
                round_decimal=round_decimal,
                expr_params=req.expr_params,
                use_default_consistency=use_default_consistency,
                **req_kwargs,
            )
            requests.append(search_request)

        hybrid_search_request = Prepare.hybrid_search_request_with_ranker(
            collection_name,
            requests,
            rerank,
            limit,
            partition_names,
            output_fields,
            round_decimal,
            use_default_consistency=use_default_consistency,
            **kwargs,
        )
        return await self._execute_hybrid_search(
            hybrid_search_request, timeout, context=context, round_decimal=round_decimal, **kwargs
        )

    @retry_on_rpc_failure()
    async def create_index(
        self,
        collection_name: str,
        field_name: str,
        params: Dict,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        index_name = kwargs.pop("index_name", Config.IndexName)

        await self.ensure_channel_ready()
        # Note: Field validation is handled by the server.
        # Client-side validation for nested fields (e.g., "chunks[text_vector]")
        # is not reliable as it only checks top-level field names.
        # Let the server perform the validation instead.
        index_param = Prepare.create_index_request(
            collection_name, field_name, params, index_name=index_name
        )

        status = await self._async_stub.CreateIndex(
            index_param, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

        index_success, fail_reason = await self.wait_for_creating_index(
            collection_name=collection_name,
            index_name=index_name,
            timeout=timeout,
            field_name=field_name,
            context=context,
            **kwargs,
        )

        if not index_success:
            raise MilvusException(message=fail_reason)

        return Status(status.code, status.reason)

    @retry_on_rpc_failure()
    async def wait_for_creating_index(
        self,
        collection_name: str,
        index_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        timestamp = await self.alloc_timestamp()
        start = time.time()
        while True:
            await asyncio.sleep(0.5)
            state, fail_reason = await self.get_index_state(
                collection_name,
                index_name,
                timeout=timeout,
                timestamp=timestamp,
                context=context,
                **kwargs,
            )
            if state == IndexState.Finished:
                return True, fail_reason
            if state == IndexState.Failed:
                return False, fail_reason
            end = time.time()
            if isinstance(timeout, int) and end - start > timeout:
                msg = (
                    f"collection {collection_name} create index {index_name} "
                    f"timeout in {timeout}s"
                )
                raise MilvusException(message=msg)

    @retry_on_rpc_failure()
    async def get_index_state(
        self,
        collection_name: str,
        index_name: str,
        timeout: Optional[float] = None,
        timestamp: Optional[int] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        request = Prepare.describe_index_request(collection_name, index_name, timestamp)
        response = await self._async_stub.DescribeIndex(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        status = response.status
        check_status(status)

        if len(response.index_descriptions) == 1:
            index_desc = response.index_descriptions[0]
            return index_desc.state, index_desc.index_state_fail_reason
        field_name = kwargs.pop("field_name", "")
        if field_name != "":
            for index_desc in response.index_descriptions:
                if index_desc.field_name == field_name:
                    return index_desc.state, index_desc.index_state_fail_reason

        raise AmbiguousIndexName(message=ExceptionsMessage.AmbiguousIndexName)

    @retry_on_rpc_failure()
    async def drop_index(
        self,
        collection_name: str,
        field_name: str,
        index_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.drop_index_request(collection_name, field_name, index_name)
        response = await self._async_stub.DropIndex(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def create_partition(
        self,
        collection_name: str,
        partition_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(
            collection_name=collection_name, partition_name=partition_name, timeout=timeout
        )
        request = Prepare.create_partition_request(collection_name, partition_name)
        response = await self._async_stub.CreatePartition(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def drop_partition(
        self,
        collection_name: str,
        partition_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(
            collection_name=collection_name, partition_name=partition_name, timeout=timeout
        )
        request = Prepare.drop_partition_request(collection_name, partition_name)

        response = await self._async_stub.DropPartition(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def load_partitions(
        self,
        collection_name: str,
        partition_names: List[str],
        replica_number: Optional[int] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(timeout=timeout)

        request = Prepare.load_partitions(
            collection_name=collection_name,
            partition_names=partition_names,
            replica_number=replica_number,
            **kwargs,
        )
        response = await self._async_stub.LoadPartitions(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

        await self.wait_for_loading_partitions(
            collection_name=collection_name,
            partition_names=partition_names,
            is_refresh=request.refresh,
            timeout=timeout,
            context=context,
            **kwargs,
        )

    @retry_on_rpc_failure()
    async def wait_for_loading_partitions(
        self,
        collection_name: str,
        partition_names: List[str],
        timeout: Optional[float] = None,
        is_refresh: bool = False,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        start = time.time()

        def can_loop(t: int) -> bool:
            return True if timeout is None else t <= (start + timeout)

        while can_loop(time.time()):
            progress = await self.get_loading_progress(
                collection_name,
                partition_names,
                timeout=timeout,
                is_refresh=is_refresh,
                context=context,
                **kwargs,
            )
            if progress >= 100:
                return
            await asyncio.sleep(Config.WaitTimeDurationWhenLoad)
        raise MilvusException(
            message=f"wait for loading partition timeout, collection: {collection_name}, partitions: {partition_names}"
        )

    @retry_on_rpc_failure()
    async def release_partitions(
        self,
        collection_name: str,
        partition_names: List[str],
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(
            collection_name=collection_name, partition_name_array=partition_names, timeout=timeout
        )
        request = Prepare.release_partitions("", collection_name, partition_names)
        response = await self._async_stub.ReleasePartitions(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def has_partition(
        self,
        collection_name: str,
        partition_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> bool:
        await self.ensure_channel_ready()
        check_pass_param(
            collection_name=collection_name, partition_name=partition_name, timeout=timeout
        )
        request = Prepare.has_partition_request(collection_name, partition_name)
        response = await self._async_stub.HasPartition(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return response.value

    @retry_on_rpc_failure()
    async def list_partitions(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> List[str]:
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.show_partitions_request(collection_name)
        response = await self._async_stub.ShowPartitions(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return list(response.partition_names)

    @retry_on_rpc_failure()
    async def get(
        self,
        collection_name: str,
        ids: List[int],
        output_fields: Optional[List[str]] = None,
        partition_names: Optional[List[str]] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        # TODO: some check
        await self.ensure_channel_ready()
        request = Prepare.retrieve_request(collection_name, ids, output_fields, partition_names)
        return await self._async_stub.Retrieve.get(
            request, timeout=timeout, metadata=_api_level_md(context)
        )

    @retry_on_rpc_failure()
    async def query(
        self,
        collection_name: str,
        expr: str,
        output_fields: Optional[List[str]] = None,
        partition_names: Optional[List[str]] = None,
        timeout: Optional[float] = None,
        strict_float32: bool = False,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        if output_fields is not None and not isinstance(output_fields, (list,)):
            raise ParamError(message="Invalid query format. 'output_fields' must be a list")

        use_default_consistency = ts_utils.construct_guarantee_ts(
            collection_name, kwargs, self.server_address, (context.get_db_name() if context else "")
        )

        request = Prepare.query_request(
            collection_name,
            expr,
            output_fields,
            partition_names,
            use_default_consistency=use_default_consistency,
            **kwargs,
        )

        response = await self._async_stub.Query(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)

        num_fields = len(response.fields_data)
        # check has fields
        if num_fields == 0:
            raise MilvusException(message="No fields returned")

        # check if all lists are of the same length
        it = iter(response.fields_data)
        num_entities = len_of(next(it))
        if not all(len_of(field_data) == num_entities for field_data in it):
            raise MilvusException(message="The length of fields data is inconsistent")

        _, dynamic_fields = entity_helper.extract_dynamic_field_from_result(response)
        keys = [field_data.field_name for field_data in response.fields_data]
        filtered_keys = [k for k in keys if k != "$meta"]
        results = [dict.fromkeys(filtered_keys) for _ in range(num_entities)]
        lazy_field_data = []
        for field_data in response.fields_data:
            lazy_extracted = entity_helper.extract_row_data_from_fields_data_v2(field_data, results)
            if lazy_extracted:
                lazy_field_data.append(field_data)

        extra_dict = get_extra_info(response.status)
        extra_dict[ITERATOR_SESSION_TS_FIELD] = response.session_ts
        return HybridExtraList(
            lazy_field_data,
            results,
            extra=extra_dict,
            dynamic_fields=dynamic_fields,
            strict_float32=strict_float32,
        )

    @retry_on_rpc_failure()
    @ignore_unimplemented(0)
    async def alloc_timestamp(
        self, timeout: Optional[float] = None, context: Optional[CallContext] = None, **kwargs
    ) -> int:
        await self.ensure_channel_ready()
        request = milvus_types.AllocTimestampRequest()
        response = await self._async_stub.AllocTimestamp(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return response.timestamp

    @retry_on_rpc_failure()
    async def alter_collection_properties(
        self,
        collection_name: str,
        properties: dict,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, properties=properties, timeout=timeout)
        request = Prepare.alter_collection_request(collection_name, properties=properties)
        status = await self._async_stub.AlterCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def drop_collection_properties(
        self,
        collection_name: str,
        property_keys: List[str],
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.alter_collection_request(collection_name, delete_keys=property_keys)
        status = await self._async_stub.AlterCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def alter_collection_field(
        self,
        collection_name: str,
        field_name: str,
        field_params: dict,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, properties=field_params, timeout=timeout)
        request = Prepare.alter_collection_field_request(
            collection_name=collection_name, field_name=field_name, field_param=field_params
        )
        status = await self._async_stub.AlterCollectionField(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def add_collection_field(
        self,
        collection_name: str,
        field_schema: FieldSchema,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.add_collection_field_request(collection_name, field_schema)
        status = await self._async_stub.AddCollectionField(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def drop_collection_function(
        self,
        collection_name: str,
        function_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.drop_collection_function_request(collection_name, function_name)

        status = await self._async_stub.DropCollectionFunction(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def add_collection_function(
        self,
        collection_name: str,
        function: Function,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.add_collection_function_request(collection_name, function)

        status = await self._async_stub.AddCollectionFunction(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def alter_collection_function(
        self,
        collection_name: str,
        function_name: str,
        function: Function,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.alter_collection_function_request(
            collection_name, function_name, function
        )

        status = await self._async_stub.AlterCollectionFunction(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def list_indexes(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.describe_index_request(collection_name, "")

        response = await self._async_stub.DescribeIndex(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        status = response.status
        if is_successful(status):
            return response.index_descriptions
        if status.code == ErrorCode.INDEX_NOT_FOUND or status.error_code == Status.INDEX_NOT_EXIST:
            return []
        raise MilvusException(status.code, status.reason, status.error_code)

    @retry_on_rpc_failure()
    async def describe_index(
        self,
        collection_name: str,
        index_name: str,
        timeout: Optional[float] = None,
        timestamp: Optional[int] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, index_name=index_name, timeout=timeout)
        request = Prepare.describe_index_request(collection_name, index_name, timestamp=timestamp)

        response = await self._async_stub.DescribeIndex(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        status = response.status
        if status.code == ErrorCode.INDEX_NOT_FOUND or status.error_code == Status.INDEX_NOT_EXIST:
            return None
        check_status(status)
        if len(response.index_descriptions) == 1:
            info_dict = {kv.key: kv.value for kv in response.index_descriptions[0].params}
            info_dict["field_name"] = response.index_descriptions[0].field_name
            info_dict["index_name"] = response.index_descriptions[0].index_name
            if info_dict.get("params"):
                info_dict["params"] = json.loads(info_dict["params"])
            info_dict["total_rows"] = response.index_descriptions[0].total_rows
            info_dict["indexed_rows"] = response.index_descriptions[0].indexed_rows
            info_dict["pending_index_rows"] = response.index_descriptions[0].pending_index_rows
            info_dict["state"] = common_pb2.IndexState.Name(response.index_descriptions[0].state)
            return info_dict

        raise AmbiguousIndexName(message=ExceptionsMessage.AmbiguousIndexName)

    @retry_on_rpc_failure()
    async def alter_index_properties(
        self,
        collection_name: str,
        index_name: str,
        properties: dict,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, index_name=index_name, timeout=timeout)
        if properties is None:
            raise ParamError(message="properties should not be None")

        request = Prepare.alter_index_properties_request(collection_name, index_name, properties)
        response = await self._async_stub.AlterIndex(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def drop_index_properties(
        self,
        collection_name: str,
        index_name: str,
        property_keys: List[str],
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, index_name=index_name, timeout=timeout)
        request = Prepare.drop_index_properties_request(
            collection_name, index_name, delete_keys=property_keys
        )
        response = await self._async_stub.AlterIndex(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def create_alias(
        self,
        collection_name: str,
        alias: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.create_alias_request(collection_name, alias)
        response = await self._async_stub.CreateAlias(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def drop_alias(
        self,
        alias: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        request = Prepare.drop_alias_request(alias)
        response = await self._async_stub.DropAlias(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def alter_alias(
        self,
        collection_name: str,
        alias: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.alter_alias_request(collection_name, alias)
        response = await self._async_stub.AlterAlias(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response)

    @retry_on_rpc_failure()
    async def describe_alias(
        self,
        alias: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(alias=alias, timeout=timeout)
        request = Prepare.describe_alias_request(alias)
        response = await self._async_stub.DescribeAlias(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        ret = {
            "alias": alias,
        }
        if response.collection:
            ret["collection_name"] = response.collection
        if response.db_name:
            ret["db_name"] = response.db_name
        return ret

    @retry_on_rpc_failure()
    async def list_aliases(
        self,
        collection_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(collection_name=collection_name, timeout=timeout)
        request = Prepare.list_aliases_request(collection_name)
        response = await self._async_stub.ListAliases(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return response.aliases

    def reset_db_name(self, db_name: str):
        """Deprecated: db_name is now passed per-request via kwargs.

        This method is kept for backward compatibility but does nothing.
        Use AsyncMilvusClient.use_database() instead.
        """

    @retry_on_rpc_failure()
    async def create_database(
        self,
        db_name: str,
        properties: Optional[dict] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(db_name=db_name, timeout=timeout)
        request = Prepare.create_database_req(db_name, properties=properties)
        status = await self._async_stub.CreateDatabase(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def drop_database(
        self,
        db_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        request = Prepare.drop_database_req(db_name)
        status = await self._async_stub.DropDatabase(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def list_database(
        self,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        check_pass_param(timeout=timeout)
        request = Prepare.list_database_req()
        response = await self._async_stub.ListDatabases(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return list(response.db_names)

    @retry_on_rpc_failure()
    async def alter_database(
        self,
        db_name: str,
        properties: dict,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        request = Prepare.alter_database_properties_req(db_name, properties)
        status = await self._async_stub.AlterDatabase(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def drop_database_properties(
        self,
        db_name: str,
        property_keys: List[str],
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        request = Prepare.drop_database_properties_req(db_name, property_keys)
        status = await self._async_stub.AlterDatabase(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(status)

    @retry_on_rpc_failure()
    async def describe_database(
        self,
        db_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        request = Prepare.describe_database_req(db_name=db_name)
        resp = await self._async_stub.DescribeDatabase(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return DatabaseInfo(resp).to_dict()

    @retry_on_rpc_failure()
    async def create_privilege_group(
        self,
        privilege_group: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.create_privilege_group_req(privilege_group)
        resp = await self._async_stub.CreatePrivilegeGroup(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def drop_privilege_group(
        self,
        privilege_group: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.drop_privilege_group_req(privilege_group)
        resp = await self._async_stub.DropPrivilegeGroup(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def list_privilege_groups(
        self,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.list_privilege_groups_req()
        resp = await self._async_stub.ListPrivilegeGroups(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return resp.privilege_groups

    @retry_on_rpc_failure()
    async def add_privileges_to_group(
        self,
        privilege_group: str,
        privileges: List[str],
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.operate_privilege_group_req(
            privilege_group, privileges, milvus_types.OperatePrivilegeGroupType.AddPrivilegesToGroup
        )
        resp = await self._async_stub.OperatePrivilegeGroup(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def remove_privileges_from_group(
        self,
        privilege_group: str,
        privileges: List[str],
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.operate_privilege_group_req(
            privilege_group,
            privileges,
            milvus_types.OperatePrivilegeGroupType.RemovePrivilegesFromGroup,
        )
        resp = await self._async_stub.OperatePrivilegeGroup(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def create_user(
        self,
        user: str,
        password: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        check_pass_param(user=user, password=password, timeout=timeout)
        req = Prepare.create_user_request(user, password)
        resp = await self._async_stub.CreateCredential(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def drop_user(
        self,
        user: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.delete_user_request(user)
        resp = await self._async_stub.DeleteCredential(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def update_password(
        self,
        user: str,
        old_password: str,
        new_password: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.update_password_request(user, old_password, new_password)
        resp = await self._async_stub.UpdateCredential(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def list_users(
        self,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.list_usernames_request()
        resp = await self._async_stub.ListCredUsers(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return resp.usernames

    @retry_on_rpc_failure()
    async def describe_user(
        self,
        username: str,
        include_role_info: bool,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.select_user_request(username, include_role_info)
        resp = await self._async_stub.SelectUser(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return resp

    @retry_on_rpc_failure()
    async def create_role(
        self,
        role_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.create_role_request(role_name)
        resp = await self._async_stub.CreateRole(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def drop_role(
        self,
        role_name: str,
        force_drop: bool = False,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.drop_role_request(role_name, force_drop=force_drop)
        resp = await self._async_stub.DropRole(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def grant_role(
        self,
        username: str,
        role_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.operate_user_role_request(
            username, role_name, milvus_types.OperateUserRoleType.AddUserToRole
        )
        resp = await self._async_stub.OperateUserRole(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def revoke_role(
        self,
        username: str,
        role_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.operate_user_role_request(
            username, role_name, milvus_types.OperateUserRoleType.RemoveUserFromRole
        )
        resp = await self._async_stub.OperateUserRole(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def describe_role(
        self,
        role_name: str,
        include_user_info: bool,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.select_role_request(role_name, include_user_info)
        resp = await self._async_stub.SelectRole(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return resp.results

    @retry_on_rpc_failure()
    async def list_roles(
        self,
        include_user_info: bool,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.select_role_request(None, include_user_info)
        resp = await self._async_stub.SelectRole(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return resp.results

    @retry_on_rpc_failure()
    async def select_grant_for_one_role(
        self,
        role_name: str,
        db_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.select_grant_request(role_name, None, None, db_name)
        resp = await self._async_stub.SelectGrant(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)

        return GrantInfo(resp.entities)

    @retry_on_rpc_failure()
    async def grant_privilege(
        self,
        role_name: str,
        object: str,
        object_name: str,
        privilege: str,
        db_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.operate_privilege_request(
            role_name,
            object,
            object_name,
            privilege,
            db_name,
            milvus_types.OperatePrivilegeType.Grant,
        )
        resp = await self._async_stub.OperatePrivilege(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def revoke_privilege(
        self,
        role_name: str,
        object: str,
        object_name: str,
        privilege: str,
        db_name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.operate_privilege_request(
            role_name,
            object,
            object_name,
            privilege,
            db_name,
            milvus_types.OperatePrivilegeType.Revoke,
        )
        resp = await self._async_stub.OperatePrivilege(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def grant_privilege_v2(
        self,
        role_name: str,
        privilege: str,
        collection_name: str,
        db_name: Optional[str] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.operate_privilege_v2_request(
            role_name,
            privilege,
            milvus_types.OperatePrivilegeType.Grant,
            db_name,
            collection_name,
        )
        resp = await self._async_stub.OperatePrivilegeV2(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def revoke_privilege_v2(
        self,
        role_name: str,
        privilege: str,
        collection_name: str,
        db_name: Optional[str] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        await self.ensure_channel_ready()
        req = Prepare.operate_privilege_v2_request(
            role_name,
            privilege,
            milvus_types.OperatePrivilegeType.Revoke,
            db_name,
            collection_name,
        )
        resp = await self._async_stub.OperatePrivilegeV2(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def create_resource_group(
        self,
        name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.create_resource_group(name, **kwargs)
        resp = await self._async_stub.CreateResourceGroup(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def drop_resource_group(
        self,
        name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.drop_resource_group(name)
        resp = await self._async_stub.DropResourceGroup(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def update_resource_groups(
        self,
        configs: Dict[str, ResourceGroupConfig],
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.update_resource_groups(configs)
        resp = await self._async_stub.UpdateResourceGroups(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def describe_resource_group(
        self,
        name: str,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.describe_resource_group(name)
        resp = await self._async_stub.DescribeResourceGroup(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return resp.resource_group

    @retry_on_rpc_failure()
    async def list_resource_groups(
        self,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.list_resource_groups()
        resp = await self._async_stub.ListResourceGroups(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)
        return list(resp.resource_groups)

    @retry_on_rpc_failure()
    async def transfer_replica(
        self,
        source: str,
        target: str,
        collection_name: str,
        num_replica: int,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.transfer_replica(source, target, collection_name, num_replica)
        resp = await self._async_stub.TransferReplica(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp)

    @retry_on_rpc_failure()
    async def get_flush_state(
        self,
        segment_ids: List[int],
        collection_name: str,
        flush_ts: int,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> bool:
        """Get the flush state for given segments."""
        await self.ensure_channel_ready()
        req = Prepare.get_flush_state_request(segment_ids, collection_name, flush_ts)
        response = await self._async_stub.GetFlushState(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)
        return response.flushed

    async def _wait_for_flushed(
        self,
        segment_ids: List[int],
        collection_name: str,
        flush_ts: int,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        """Wait for segments to be flushed."""
        flush_ret = False
        start = time.time()
        while not flush_ret:
            flush_ret = await self.get_flush_state(
                segment_ids, collection_name, flush_ts, timeout, context=context, **kwargs
            )
            end = time.time()
            if timeout is not None and end - start > timeout:
                raise MilvusException(
                    message=f"wait for flush timeout, collection: {collection_name}, flusht_ts: {flush_ts}"
                )

            if not flush_ret:
                await asyncio.sleep(0.5)

    @retry_on_rpc_failure()
    async def flush(
        self,
        collection_names: List[str],
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        if collection_names in (None, []) or not isinstance(collection_names, list):
            raise ParamError(message="Collection name list can not be None or empty")

        check_pass_param(timeout=timeout)
        for name in collection_names:
            check_pass_param(collection_name=name)

        req = Prepare.flush_param(collection_names)
        response = await self._async_stub.Flush(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)

        # Wait for all segments to be flushed in parallel
        if collection_names:
            await asyncio.gather(
                *(
                    self._wait_for_flushed(
                        response.coll_segIDs[collection_name].data,
                        collection_name,
                        response.coll_flush_ts[collection_name],
                        timeout=timeout,
                        context=context,
                    )
                    for collection_name in collection_names
                )
            )

        return response

    @retry_on_rpc_failure()
    async def compact(
        self,
        collection_name: str,
        is_clustering: Optional[bool] = False,
        is_l0: Optional[bool] = False,
        target_size: Optional[int] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ) -> int:

        request = Prepare.describe_collection_request(collection_name)
        response = await self._async_stub.DescribeCollection(
            request, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)

        req = Prepare.manual_compaction(
            collection_name=collection_name,
            is_clustering=is_clustering,
            is_l0=is_l0,
            collection_id=response.collectionID,
            target_size=target_size,
        )
        response = await self._async_stub.ManualCompaction(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)

        return response.compactionID

    @retry_on_rpc_failure()
    async def get_compaction_state(
        self,
        compaction_id: int,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.get_compaction_state(compaction_id)
        response = await self._async_stub.GetCompactionState(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(response.status)

        return CompactionState(
            compaction_id,
            State.new(response.state),
            response.executingPlanNo,
            response.timeoutPlanNo,
            response.completedPlanNo,
        )

    @retry_on_rpc_failure()
    async def run_analyzer(
        self,
        texts: Union[str, List[str]],
        analyzer_params: Optional[Union[str, Dict]] = None,
        with_hash: bool = False,
        with_detail: bool = False,
        collection_name: Optional[str] = None,
        field_name: Optional[str] = None,
        analyzer_names: Optional[Union[str, List[str]]] = None,
        timeout: Optional[float] = None,
        context: Optional[CallContext] = None,
        **kwargs,
    ):
        req = Prepare.run_analyzer(
            texts,
            analyzer_params=analyzer_params,
            with_hash=with_hash,
            with_detail=with_detail,
            collection_name=collection_name,
            field_name=field_name,
            analyzer_names=analyzer_names,
        )
        resp = await self._async_stub.RunAnalyzer(
            req, timeout=timeout, metadata=_api_level_md(context)
        )
        check_status(resp.status)

        if isinstance(texts, str):
            return AnalyzeResult(resp.results[0], with_hash, with_detail)
        return [AnalyzeResult(result, with_hash, with_detail) for result in resp.results]
