"""Async session wrapper for MXAccess Gateway commands."""

from __future__ import annotations

from collections.abc import AsyncIterator, Sequence

from .auth import redact_secret
from .errors import MalformedReplyError, MxGatewayError, ensure_mxaccess_success
from .events import ReplayGap
from .generated import mxaccess_gateway_pb2 as pb
from .values import MxValueInput, to_mx_value

MAX_BULK_ITEMS = 1000


class Session:
    """A single gateway-backed MXAccess session."""

    def __init__(
        self,
        *,
        client: "GatewayClient",
        session_id: str,
        open_reply: pb.OpenSessionReply | None = None,
    ) -> None:
        """Initialize a session bound to a client and gateway session id."""
        self.client = client
        self.session_id = session_id
        self.open_reply = open_reply
        self._closed = False

    async def __aenter__(self) -> "Session":
        """Return self to support ``async with`` usage."""
        return self

    async def __aexit__(self, *_exc_info: object) -> None:
        """Close the session when leaving an ``async with`` block."""
        await self.close()

    async def close(self, *, client_correlation_id: str = "") -> pb.CloseSessionReply:
        """Close the gateway session. Repeated calls return a local closed reply."""

        if self._closed:
            return pb.CloseSessionReply(
                session_id=self.session_id,
                final_state=pb.SESSION_STATE_CLOSED,
                protocol_status=pb.ProtocolStatus(code=pb.PROTOCOL_STATUS_CODE_OK),
            )

        reply = await self.client.close_session_raw(
            pb.CloseSessionRequest(
                session_id=self.session_id,
                client_correlation_id=client_correlation_id,
            ),
        )
        self._closed = True
        return reply

    async def invoke(self, command: pb.MxCommand, *, correlation_id: str = "") -> pb.MxCommandReply:
        """Invoke a raw command and enforce gateway and MXAccess success."""

        reply = await self.invoke_raw(command, correlation_id=correlation_id)
        return ensure_mxaccess_success("invoke", reply)

    async def invoke_raw(
        self,
        command: pb.MxCommand,
        *,
        correlation_id: str = "",
    ) -> pb.MxCommandReply:
        """Invoke a raw command and preserve the raw reply."""

        return await self.client.invoke_raw(
            pb.MxCommandRequest(
                session_id=self.session_id,
                client_correlation_id=correlation_id,
                command=command,
            ),
        )

    async def register(self, client_name: str, *, correlation_id: str = "") -> int:
        """Invoke MXAccess `Register` and return the new `ServerHandle`."""
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_REGISTER,
                register=pb.RegisterCommand(client_name=client_name),
            ),
            correlation_id=correlation_id,
        )
        return reply.register.server_handle

    async def unregister(self, server_handle: int, *, correlation_id: str = "") -> None:
        """Invoke MXAccess `Unregister` for a previously registered `ServerHandle`."""
        await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_UNREGISTER,
                unregister=pb.UnregisterCommand(server_handle=server_handle),
            ),
            correlation_id=correlation_id,
        )

    async def remove_item(
        self,
        server_handle: int,
        item_handle: int,
        *,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `RemoveItem` for the given `ItemHandle`."""
        await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_REMOVE_ITEM,
                remove_item=pb.RemoveItemCommand(
                    server_handle=server_handle,
                    item_handle=item_handle,
                ),
            ),
            correlation_id=correlation_id,
        )

    async def add_item(
        self,
        server_handle: int,
        item_definition: str,
        *,
        correlation_id: str = "",
    ) -> int:
        """Invoke MXAccess `AddItem` and return the new `ItemHandle`."""
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ADD_ITEM,
                add_item=pb.AddItemCommand(
                    server_handle=server_handle,
                    item_definition=item_definition,
                ),
            ),
            correlation_id=correlation_id,
        )
        return reply.add_item.item_handle

    async def add_item2(
        self,
        server_handle: int,
        item_definition: str,
        item_context: str,
        *,
        correlation_id: str = "",
    ) -> int:
        """Invoke MXAccess `AddItem2` with item context and return the new `ItemHandle`."""
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ADD_ITEM2,
                add_item2=pb.AddItem2Command(
                    server_handle=server_handle,
                    item_definition=item_definition,
                    item_context=item_context,
                ),
            ),
            correlation_id=correlation_id,
        )
        return reply.add_item2.item_handle

    async def advise(
        self,
        server_handle: int,
        item_handle: int,
        *,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `Advise` to subscribe an existing `ItemHandle` to events."""
        await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ADVISE,
                advise=pb.AdviseCommand(
                    server_handle=server_handle,
                    item_handle=item_handle,
                ),
            ),
            correlation_id=correlation_id,
        )

    async def unadvise(
        self,
        server_handle: int,
        item_handle: int,
        *,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `UnAdvise` to stop event delivery for an `ItemHandle`."""
        await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_UN_ADVISE,
                un_advise=pb.UnAdviseCommand(
                    server_handle=server_handle,
                    item_handle=item_handle,
                ),
            ),
            correlation_id=correlation_id,
        )

    async def add_item_bulk(
        self,
        server_handle: int,
        tag_addresses: Sequence[str],
        *,
        correlation_id: str = "",
    ) -> list[pb.SubscribeResult]:
        """Invoke MXAccess `AddItemBulk` and return one result per tag address."""
        if tag_addresses is None:
            raise TypeError("tag_addresses is required")
        _ensure_bulk_size("tag_addresses", len(tag_addresses))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ADD_ITEM_BULK,
                add_item_bulk=pb.AddItemBulkCommand(
                    server_handle=server_handle,
                    tag_addresses=tag_addresses,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.add_item_bulk.results)

    async def advise_item_bulk(
        self,
        server_handle: int,
        item_handles: Sequence[int],
        *,
        correlation_id: str = "",
    ) -> list[pb.SubscribeResult]:
        """Invoke MXAccess `AdviseItemBulk` and return one result per item handle."""
        if item_handles is None:
            raise TypeError("item_handles is required")
        _ensure_bulk_size("item_handles", len(item_handles))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ADVISE_ITEM_BULK,
                advise_item_bulk=pb.AdviseItemBulkCommand(
                    server_handle=server_handle,
                    item_handles=item_handles,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.advise_item_bulk.results)

    async def remove_item_bulk(
        self,
        server_handle: int,
        item_handles: Sequence[int],
        *,
        correlation_id: str = "",
    ) -> list[pb.SubscribeResult]:
        """Invoke MXAccess `RemoveItemBulk` and return one result per item handle."""
        if item_handles is None:
            raise TypeError("item_handles is required")
        _ensure_bulk_size("item_handles", len(item_handles))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_REMOVE_ITEM_BULK,
                remove_item_bulk=pb.RemoveItemBulkCommand(
                    server_handle=server_handle,
                    item_handles=item_handles,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.remove_item_bulk.results)

    async def unadvise_item_bulk(
        self,
        server_handle: int,
        item_handles: Sequence[int],
        *,
        correlation_id: str = "",
    ) -> list[pb.SubscribeResult]:
        """Invoke MXAccess `UnAdviseItemBulk` and return one result per item handle."""
        if item_handles is None:
            raise TypeError("item_handles is required")
        _ensure_bulk_size("item_handles", len(item_handles))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_UN_ADVISE_ITEM_BULK,
                un_advise_item_bulk=pb.UnAdviseItemBulkCommand(
                    server_handle=server_handle,
                    item_handles=item_handles,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.un_advise_item_bulk.results)

    async def subscribe_bulk(
        self,
        server_handle: int,
        tag_addresses: Sequence[str],
        *,
        correlation_id: str = "",
    ) -> list[pb.SubscribeResult]:
        """Invoke MXAccess `SubscribeBulk` and return one result per tag address."""
        if tag_addresses is None:
            raise TypeError("tag_addresses is required")
        _ensure_bulk_size("tag_addresses", len(tag_addresses))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_SUBSCRIBE_BULK,
                subscribe_bulk=pb.SubscribeBulkCommand(
                    server_handle=server_handle,
                    tag_addresses=tag_addresses,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.subscribe_bulk.results)

    async def unsubscribe_bulk(
        self,
        server_handle: int,
        item_handles: Sequence[int],
        *,
        correlation_id: str = "",
    ) -> list[pb.SubscribeResult]:
        """Invoke MXAccess `UnsubscribeBulk` and return one result per item handle."""
        if item_handles is None:
            raise TypeError("item_handles is required")
        _ensure_bulk_size("item_handles", len(item_handles))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_UNSUBSCRIBE_BULK,
                unsubscribe_bulk=pb.UnsubscribeBulkCommand(
                    server_handle=server_handle,
                    item_handles=item_handles,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.unsubscribe_bulk.results)

    async def write_bulk(
        self,
        server_handle: int,
        entries: Sequence[pb.WriteBulkEntry],
        *,
        correlation_id: str = "",
    ) -> list[pb.BulkWriteResult]:
        """Invoke MXAccess `WriteBulk` and return one BulkWriteResult per entry.

        Per-entry MXAccess failures appear as results with ``was_successful = False``
        and a populated ``error_message`` / ``hresult``; this method does not raise
        on per-entry failure, mirroring the existing add/advise bulk surface.
        """
        if entries is None:
            raise TypeError("entries is required")
        _ensure_bulk_size("entries", len(entries))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_WRITE_BULK,
                write_bulk=pb.WriteBulkCommand(
                    server_handle=server_handle,
                    entries=entries,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.write_bulk.results)

    async def write2_bulk(
        self,
        server_handle: int,
        entries: Sequence[pb.Write2BulkEntry],
        *,
        correlation_id: str = "",
    ) -> list[pb.BulkWriteResult]:
        """Invoke MXAccess `Write2Bulk` (timestamped) and return per-entry results."""
        if entries is None:
            raise TypeError("entries is required")
        _ensure_bulk_size("entries", len(entries))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_WRITE2_BULK,
                write2_bulk=pb.Write2BulkCommand(
                    server_handle=server_handle,
                    entries=entries,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.write2_bulk.results)

    async def write_secured_bulk(
        self,
        server_handle: int,
        entries: Sequence[pb.WriteSecuredBulkEntry],
        *,
        correlation_id: str = "",
    ) -> list[pb.BulkWriteResult]:
        """Invoke MXAccess `WriteSecuredBulk` — credential-sensitive values must not be logged."""
        if entries is None:
            raise TypeError("entries is required")
        _ensure_bulk_size("entries", len(entries))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_WRITE_SECURED_BULK,
                write_secured_bulk=pb.WriteSecuredBulkCommand(
                    server_handle=server_handle,
                    entries=entries,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.write_secured_bulk.results)

    async def write_secured2_bulk(
        self,
        server_handle: int,
        entries: Sequence[pb.WriteSecured2BulkEntry],
        *,
        correlation_id: str = "",
    ) -> list[pb.BulkWriteResult]:
        """Invoke MXAccess `WriteSecured2Bulk` (timestamped + verified)."""
        if entries is None:
            raise TypeError("entries is required")
        _ensure_bulk_size("entries", len(entries))
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_WRITE_SECURED2_BULK,
                write_secured2_bulk=pb.WriteSecured2BulkCommand(
                    server_handle=server_handle,
                    entries=entries,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.write_secured2_bulk.results)

    async def read_bulk(
        self,
        server_handle: int,
        tag_addresses: Sequence[str],
        *,
        timeout_ms: int = 0,
        correlation_id: str = "",
    ) -> list[pb.BulkReadResult]:
        """Invoke `ReadBulk` — snapshot the current value of each requested tag.

        MXAccess COM has no synchronous read; the worker returns the cached
        ``OnDataChange`` value for any tag that is already advised (``was_cached =
        True``) without modifying the existing subscription, and falls back to
        a full AddItem + Advise + wait + UnAdvise + RemoveItem snapshot lifecycle
        otherwise. ``timeout_ms`` bounds the per-tag wait in the snapshot case;
        pass ``0`` to use the worker default (1000 ms).
        """
        if tag_addresses is None:
            raise TypeError("tag_addresses is required")
        _ensure_bulk_size("tag_addresses", len(tag_addresses))
        if timeout_ms < 0:
            raise ValueError("timeout_ms must be non-negative")
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_READ_BULK,
                read_bulk=pb.ReadBulkCommand(
                    server_handle=server_handle,
                    tag_addresses=tag_addresses,
                    timeout_ms=timeout_ms,
                ),
            ),
            correlation_id=correlation_id,
        )
        return list(reply.read_bulk.results)

    async def write(
        self,
        server_handle: int,
        item_handle: int,
        value: MxValueInput,
        *,
        user_id: int = 0,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `Write` for an `ItemHandle` with the converted value."""
        await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_WRITE,
                write=pb.WriteCommand(
                    server_handle=server_handle,
                    item_handle=item_handle,
                    value=to_mx_value(value),
                    user_id=user_id,
                ),
            ),
            correlation_id=correlation_id,
        )

    async def write_array_elements(
        self,
        server_handle: int,
        item_handle: int,
        element_data_type: "pb.MxDataType.ValueType",
        total_length: int,
        elements: dict[int, MxValueInput],
        *,
        user_id: int = 0,
        correlation_id: str = "",
    ) -> None:
        """Write a partial array by specifying only the indices you want to set.

        The gateway expands the sparse representation into a full ``total_length``
        array before forwarding the write to MXAccess.  Indices not listed in
        *elements* are filled with the type default for *element_data_type* (0,
        False, empty string, Unix epoch for timestamps, etc.).  The previous
        value at those positions is **not** preserved — this is a full array
        replacement, not a patch.

        Args:
            server_handle: Handle returned by :meth:`register`.
            item_handle: Handle returned by :meth:`add_item`.
            element_data_type: ``pb.MX_DATA_TYPE_*`` enum value for the scalar
                element type of the target array attribute.
            total_length: Total number of elements in the written array.  Must
                be > 0 and large enough to contain every index in *elements*.
                Both *total_length* and all keys in *elements* must be
                non-negative; the gateway rejects negative or out-of-range
                values with ``InvalidArgument`` (the proto fields are
                ``uint32``).
            elements: Mapping of zero-based element index to scalar value.
                Values are converted with :func:`~zb_mom_ww_mxgateway.values.to_mx_value`.
            user_id: Galaxy user id to stamp on the write (requires a prior
                supervisory advise to take effect — see README).
            correlation_id: Optional client-supplied correlation token echoed
                in the command reply.
        """
        sparse = pb.MxSparseArray(
            element_data_type=element_data_type,
            total_length=total_length,
            elements=[
                pb.MxSparseElement(index=idx, value=to_mx_value(val))
                for idx, val in elements.items()
            ],
        )
        await self.write(
            server_handle,
            item_handle,
            pb.MxValue(sparse_array_value=sparse),
            user_id=user_id,
            correlation_id=correlation_id,
        )

    async def write2(
        self,
        server_handle: int,
        item_handle: int,
        value: MxValueInput,
        timestamp_value: MxValueInput,
        *,
        user_id: int = 0,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `Write2` with both a value and a client-supplied timestamp."""
        await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_WRITE2,
                write2=pb.Write2Command(
                    server_handle=server_handle,
                    item_handle=item_handle,
                    value=to_mx_value(value),
                    timestamp_value=to_mx_value(timestamp_value),
                    user_id=user_id,
                ),
            ),
            correlation_id=correlation_id,
        )

    async def _invoke_redacted(
        self,
        command: pb.MxCommand,
        *,
        correlation_id: str,
        secrets: Sequence[str | None],
    ) -> pb.MxCommandReply:
        """Invoke a command whose request carries credential-sensitive data.

        Runs the same gateway + MXAccess validation as :meth:`invoke`, but scrubs
        the supplied secret substrings from any surfaced error message before it
        propagates. MXAccess parity is preserved — the native failure is still
        raised as :class:`~zb_mom_ww_mxgateway.errors.MxAccessError`; only the
        credential text is removed from the message so it can never reach logs.
        """

        try:
            return await self.invoke(command, correlation_id=correlation_id)
        except MxGatewayError as error:
            _redact_error(error, secrets)
            raise

    async def advise_supervisory(
        self,
        server_handle: int,
        item_handle: int,
        *,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `AdviseSupervisory` for an `ItemHandle`.

        Supervisory advise is the prerequisite for user-attributed and secured
        writes: it must be established (typically after :meth:`authenticate_user`)
        before a ``user_id``-bearing :meth:`write` or a :meth:`write_secured`
        takes effect.
        """
        await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ADVISE_SUPERVISORY,
                advise_supervisory=pb.AdviseSupervisoryCommand(
                    server_handle=server_handle,
                    item_handle=item_handle,
                ),
            ),
            correlation_id=correlation_id,
        )

    async def write_secured(
        self,
        server_handle: int,
        item_handle: int,
        value: MxValueInput,
        *,
        current_user_id: int = 0,
        verifier_user_id: int = 0,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `WriteSecured` — a signed/verified write.

        The written *value* is credential-sensitive and is scrubbed from any
        surfaced error message (never logged). MXAccess parity is the contract:
        ``WriteSecured`` failing before a prior :meth:`authenticate_user` +
        :meth:`advise_supervisory`, or before a value-bearing body, is the native
        behaviour and is surfaced as-is — it is not "fixed".
        """
        await self._invoke_redacted(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_WRITE_SECURED,
                write_secured=pb.WriteSecuredCommand(
                    server_handle=server_handle,
                    item_handle=item_handle,
                    current_user_id=current_user_id,
                    verifier_user_id=verifier_user_id,
                    value=to_mx_value(value),
                ),
            ),
            correlation_id=correlation_id,
            secrets=_value_secrets(value),
        )

    async def write_secured2(
        self,
        server_handle: int,
        item_handle: int,
        value: MxValueInput,
        timestamp_value: MxValueInput,
        *,
        current_user_id: int = 0,
        verifier_user_id: int = 0,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `WriteSecured2` — a signed/verified, timestamped write.

        Like :meth:`write_secured` but also stamps a client-supplied timestamp.
        The written *value* is credential-sensitive and is scrubbed from any
        surfaced error message. Native pre-condition failures are surfaced as-is.
        """
        await self._invoke_redacted(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_WRITE_SECURED2,
                write_secured2=pb.WriteSecured2Command(
                    server_handle=server_handle,
                    item_handle=item_handle,
                    current_user_id=current_user_id,
                    verifier_user_id=verifier_user_id,
                    value=to_mx_value(value),
                    timestamp_value=to_mx_value(timestamp_value),
                ),
            ),
            correlation_id=correlation_id,
            secrets=_value_secrets(value),
        )

    async def authenticate_user(
        self,
        server_handle: int,
        verify_user: str,
        verify_user_password: str,
        *,
        correlation_id: str = "",
    ) -> int:
        """Invoke MXAccess `AuthenticateUser` and return the resolved Galaxy user id.

        *verify_user_password* is a raw MXAccess credential: it is never logged
        and is scrubbed from any surfaced error message. A native authentication
        failure is surfaced as :class:`~zb_mom_ww_mxgateway.errors.MxAccessError`
        with the credential removed.
        """
        reply = await self._invoke_redacted(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_AUTHENTICATE_USER,
                authenticate_user=pb.AuthenticateUserCommand(
                    server_handle=server_handle,
                    verify_user=verify_user,
                    verify_user_password=verify_user_password,
                ),
            ),
            correlation_id=correlation_id,
            secrets=[verify_user_password],
        )
        if reply.HasField("authenticate_user"):
            return reply.authenticate_user.user_id
        if reply.HasField("return_value") and reply.return_value.WhichOneof("kind") == "int32_value":
            return reply.return_value.int32_value
        raise MalformedReplyError(
            "authenticate_user returned a malformed reply: OK reply carried "
            "neither the typed payload nor an int32 return_value",
            raw_reply=reply,
        )

    async def archestra_user_to_id(
        self,
        server_handle: int,
        user_id_guid: str,
        *,
        correlation_id: str = "",
    ) -> int:
        """Invoke MXAccess `ArchestrAUserToId` and return the resolved user id."""
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ARCHESTRA_USER_TO_ID,
                archestra_user_to_id=pb.ArchestrAUserToIdCommand(
                    server_handle=server_handle,
                    user_id_guid=user_id_guid,
                ),
            ),
            correlation_id=correlation_id,
        )
        if reply.HasField("archestra_user_to_id"):
            return reply.archestra_user_to_id.user_id
        if reply.HasField("return_value") and reply.return_value.WhichOneof("kind") == "int32_value":
            return reply.return_value.int32_value
        raise MalformedReplyError(
            "archestra_user_to_id returned a malformed reply: OK reply carried "
            "neither the typed payload nor an int32 return_value",
            raw_reply=reply,
        )

    async def add_buffered_item(
        self,
        server_handle: int,
        item_definition: str,
        item_context: str = "",
        *,
        correlation_id: str = "",
    ) -> int:
        """Invoke MXAccess `AddBufferedItem` and return the new `ItemHandle`."""
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ADD_BUFFERED_ITEM,
                add_buffered_item=pb.AddBufferedItemCommand(
                    server_handle=server_handle,
                    item_definition=item_definition,
                    item_context=item_context,
                ),
            ),
            correlation_id=correlation_id,
        )
        if reply.HasField("add_buffered_item"):
            return reply.add_buffered_item.item_handle
        if reply.HasField("return_value") and reply.return_value.WhichOneof("kind") == "int32_value":
            return reply.return_value.int32_value
        raise MalformedReplyError(
            "add_buffered_item returned a malformed reply: OK reply carried "
            "neither the typed payload nor an int32 return_value",
            raw_reply=reply,
        )

    async def set_buffered_update_interval(
        self,
        server_handle: int,
        update_interval_milliseconds: int,
        *,
        correlation_id: str = "",
    ) -> None:
        """Invoke MXAccess `SetBufferedUpdateInterval` for the server handle."""
        await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_SET_BUFFERED_UPDATE_INTERVAL,
                set_buffered_update_interval=pb.SetBufferedUpdateIntervalCommand(
                    server_handle=server_handle,
                    update_interval_milliseconds=update_interval_milliseconds,
                ),
            ),
            correlation_id=correlation_id,
        )

    async def suspend(
        self,
        server_handle: int,
        item_handle: int,
        *,
        correlation_id: str = "",
    ) -> pb.MxStatusProxy:
        """Invoke MXAccess `Suspend` for an `ItemHandle` and return its status."""
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_SUSPEND,
                suspend=pb.SuspendCommand(
                    server_handle=server_handle,
                    item_handle=item_handle,
                ),
            ),
            correlation_id=correlation_id,
        )
        return reply.suspend.status

    async def activate(
        self,
        server_handle: int,
        item_handle: int,
        *,
        correlation_id: str = "",
    ) -> pb.MxStatusProxy:
        """Invoke MXAccess `Activate` for a suspended `ItemHandle` and return its status."""
        reply = await self.invoke(
            pb.MxCommand(
                kind=pb.MX_COMMAND_KIND_ACTIVATE,
                activate=pb.ActivateCommand(
                    server_handle=server_handle,
                    item_handle=item_handle,
                ),
            ),
            correlation_id=correlation_id,
        )
        return reply.activate.status

    def stream_events(
        self,
        *,
        after_worker_sequence: int = 0,
    ) -> AsyncIterator[pb.MxEvent | ReplayGap]:
        """Return an async iterator over this session's `MxEvent` stream.

        Each yielded item is either a normal :class:`~...MxEvent` or a
        :class:`ReplayGap`. Branch on ``isinstance(item, ReplayGap)`` — a gap is
        never delivered as an ``MxEvent`` so it cannot be mistaken for a real
        MXAccess event.

        Pass a non-zero *after_worker_sequence* to resume a previously observed
        stream. If that cursor predates the oldest event the gateway still
        retains in its replay ring, the stream opens with a single
        :class:`ReplayGap` sentinel (events in the gap were dropped and cannot be
        replayed), then continues with normal events. On a gap, discard locally
        cached state, re-snapshot, and — to resume without another gap —
        reconnect with ``after_worker_sequence = gap.resume_after_worker_sequence``.
        See :class:`ReplayGap` for the full semantics. The underlying protobuf
        stream is available raw via ``GatewayClient.stream_events_raw``.
        """
        raw = self.client.stream_events_raw(
            pb.StreamEventsRequest(
                session_id=self.session_id,
                after_worker_sequence=after_worker_sequence,
            ),
        )
        return _surface_replay_gaps(raw)


async def _surface_replay_gaps(
    raw: AsyncIterator[pb.MxEvent],
) -> AsyncIterator[pb.MxEvent | ReplayGap]:
    """Map the gateway's ``replay_gap`` sentinel event to a typed :class:`ReplayGap`.

    Normal events pass through unchanged. The sentinel (``replay_gap`` set,
    ``family`` unspecified, body unset) is converted to a distinct
    :class:`ReplayGap` so a consumer can branch on it without inspecting proto
    presence, and is never yielded as an ``MxEvent``. The sentinel is forwarded
    faithfully — it is neither dropped nor turned into a normal event.

    Closing this generator (``aclose``) propagates to *raw* so the underlying
    gRPC call is cancelled, preserving the raw stream's cancel-on-stop contract.
    """
    try:
        async for event in raw:
            if event.HasField("replay_gap"):
                yield ReplayGap.from_proto(event.replay_gap)
            else:
                yield event
    finally:
        aclose = getattr(raw, "aclose", None)
        if aclose is not None:
            await aclose()


def _ensure_bulk_size(name: str, count: int) -> None:
    if count > MAX_BULK_ITEMS:
        raise ValueError(f"{name} bulk commands are limited to {MAX_BULK_ITEMS} item(s)")


def _value_secrets(value: MxValueInput) -> list[str]:
    """Return the redaction candidate strings for a credential-sensitive write value.

    Secured-write payloads may carry a password or other secret. Only textual
    values can appear verbatim in a surfaced error, so a string value (or a
    UTF-8-decodable ``bytes`` value) is returned for scrubbing; other value
    kinds have no verbatim text form to leak.
    """
    if isinstance(value, str):
        return [value] if value else []
    if isinstance(value, bytes):
        try:
            decoded = value.decode("utf-8")
        except UnicodeDecodeError:
            return []
        return [decoded] if decoded else []
    return []


def _redact_error(error: MxGatewayError, secrets: Sequence[str | None]) -> None:
    """Scrub secret substrings from a raised error's message and reply in place.

    Rewrites ``error.args[0]`` (the message returned by ``str(error)``) through
    the shared :func:`~zb_mom_ww_mxgateway.auth.redact_secret` seam so credential
    text can never reach logs or be re-raised to a caller.

    A misbehaving MXAccess provider can echo the client-supplied credential back
    verbatim in its failure diagnostics, so ``error.raw_reply`` (the protobuf
    reply) can carry the secret in ``protocol_status.message``,
    ``diagnostic_message``, and each ``statuses[].diagnostic_text``. A logger
    dumping those structured fields would reintroduce the leak the message scrub
    closes. When there is a secret to scrub and a reply is attached, this rebinds
    ``error.raw_reply`` to a scrubbed deep copy so the raised exception carries no
    credential text on any surface. The clone leaves the original reply untouched.
    """
    scrubbed = [secret for secret in secrets if secret]
    if not scrubbed:
        return
    if error.args and isinstance(error.args[0], str):
        error.args = (redact_secret(error.args[0], scrubbed), *error.args[1:])
    if error.raw_reply is not None:
        error.raw_reply = _redact_reply(error.raw_reply, scrubbed)


def _redact_reply(reply: pb.MxCommandReply, secrets: Sequence[str]) -> pb.MxCommandReply:
    """Return a deep copy of *reply* with credential text scrubbed from diagnostics.

    Operates on a clone so the caller's original reply object is never mutated.
    Only the free-text diagnostic fields that can echo a client-supplied secret
    are scrubbed; the structured/enum fields the gateway itself sets are left as-is.
    """
    clone = type(reply)()
    clone.CopyFrom(reply)
    if clone.protocol_status.message:
        clone.protocol_status.message = redact_secret(clone.protocol_status.message, secrets)
    if clone.diagnostic_message:
        clone.diagnostic_message = redact_secret(clone.diagnostic_message, secrets)
    for status in clone.statuses:
        if status.diagnostic_text:
            status.diagnostic_text = redact_secret(status.diagnostic_text, secrets)
    return clone


from .client import GatewayClient  # noqa: E402
