"""
StreamManager for bidirectional gRPC streaming.

This module provides thread-safe bidirectional stream management for
low-latency tensor operations, enabling pipelining of requests without
waiting for individual responses.

All operations (execute_aten, copy_tensor, update_tensor, get_tensor, etc.)
go through a single ordered stream to ensure proper sequencing.
"""

from __future__ import annotations

import asyncio
import collections
import gc
import logging
import os
import struct
import sys
import threading
import time
from dataclasses import dataclass
from typing import Optional

from skytorch.torch.profiler import PROFILING_ENABLED

try:
    import grpc
except ImportError as e:
    raise ImportError(f"grpcio package is required: {e}")

try:
    from skytorch.torch.server import service_pb2
    from skytorch.torch.server import service_pb2_grpc
except ImportError:
    raise ImportError("Generated gRPC code not found. Run hack/gen-grpc-proto.sh first.")

from grpc.aio._typing import MetadataType

logger = logging.getLogger(__name__)

# Chunk size for large tensor payloads (1MB, matching serialization module)
STREAM_CHUNK_SIZE = 1024 * 1024


def _serialize_tensor_metadata_binary(meta) -> bytes:
    """Serialize a TensorMetadata proto to the ATen binary tensor metadata format.

    Format: [tensor_id:u64][ndim:u8][shape:i64×ndim][stride:i64×ndim]
            [dtype_len:u8][dtype:bytes][storage_offset:i64][nbytes:i64]
            [device_type_len:u8][device_type:bytes][device_index:i32]
            [has_tensor_ref:u8][tensor_ref:u64 if has]
    """
    shape = list(meta.shape)
    stride = list(meta.stride) if meta.stride else [0] * len(shape)
    ndim = len(shape)
    dtype_bytes = meta.dtype.encode("utf-8")
    device_type_bytes = meta.device_type.encode("utf-8")
    has_ref = meta.HasField("tensor_ref")

    parts = [
        struct.pack(f"<QB{ndim}q{ndim}q", meta.tensor_id, ndim, *shape, *stride),
        struct.pack(f"<B{len(dtype_bytes)}s", len(dtype_bytes), dtype_bytes),
        struct.pack("<qq", meta.storage_offset, meta.nbytes),
        struct.pack(
            f"<B{len(device_type_bytes)}si",
            len(device_type_bytes),
            device_type_bytes,
            meta.device_index,
        ),
    ]
    if has_ref:
        parts.append(struct.pack("<BQ", 1, meta.tensor_ref))
    else:
        parts.append(b"\x00")
    return b"".join(parts)


@dataclass
class PendingRequest:
    """Tracks a pending request awaiting response."""

    future: Optional[asyncio.Future]
    request_type: str


class StreamManager:
    """
    Bidirectional stream manager for pipelined operations.

    All tensor operations go through a single ordered stream to ensure
    proper sequencing. Responses are matched to requests by order (FIFO).

    Key features:
    - FIFO ordering of requests and responses
    - Efficient async queue (no busy-wait polling)
    - Automatic stream lifecycle management

    Note: All operations must be called from the asyncio event loop thread.
    """

    def __init__(
        self,
        stub: service_pb2_grpc.ServiceStub,
        metadata: Optional[MetadataType] = None,
    ):
        """
        Initialize the stream manager.

        Args:
            stub: gRPC service stub
            metadata: Optional metadata for requests
        """
        self._stub = stub
        self._metadata = metadata

        # Pending requests - FIFO queue of futures matched by order
        self._pending: collections.deque[PendingRequest] = collections.deque()

        # Stream state
        self._request_queue: Optional[asyncio.Queue] = None
        self._sender_task: Optional[asyncio.Task] = None
        self._receiver_task: Optional[asyncio.Task] = None
        self._stream_call: Optional[grpc.aio.StreamStreamCall] = None
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._started = False
        self._closing = False

        # Deferred error from fire-and-forget operations
        self._last_error: Optional[Exception] = None

        # Batch buffer for execute_aten operations
        self._batch_buffer: list[service_pb2.ExecuteAtenRequest] = []
        self._flush_scheduled: bool = False
        self._batch_flush_timer = None

        # Batch buffer for raw binary execute_aten operations (C++ fast path)
        self._raw_batch_buffer: list[bytes] = []
        self._raw_flush_scheduled: bool = False
        self._raw_flush_timer = None

        # Main-thread submission buffer: collects ops from all submit methods
        # under a single lock, preserving FIFO ordering while reducing
        # call_soon_threadsafe calls by ~98%. Each entry is either:
        #   ("raw", bytes)    — raw binary execute_aten
        #   ("aten", ExecuteAtenRequest) — protobuf execute_aten
        #   ("req", StreamRequest) — non-batched request (copy, update, register)
        #   ("ff", list[int])      — deferred delete tensor IDs
        self._mt_ops: list[tuple] = []
        self._mt_lock = threading.RLock()
        self._mt_wake_pending: bool = False

        # Deferred delete tensor IDs: accumulated from fire-and-forget delete
        # requests that preserve FIFO ordering but don't force a batch flush.
        # They piggyback on the next batch flush instead of triggering one.
        # IDs are batched into a single DeleteTensorsRequest when flushed,
        # reducing per-message gRPC overhead.
        self._deferred_delete_ids: list[int] = []
        self._deferred_flush_scheduled: bool = False
        self._deferred_flush_timer = None

        # Shutdown signaling
        self._shutdown_event: Optional[asyncio.Event] = None

        # C++ raw submit buffer integration
        self._cpp_submit_available: bool = False

        # Free-threaded Python uses biased reference counting: decrefs on
        # non-owning threads are deferred and only reconciled during GC.
        # When deletes accumulate, we call gc.collect(0) to reconcile
        # biased refcounts — this may free additional objects whose
        # decrefs were deferred by non-owning threads (e.g., the
        # AsyncStreamer worker thread).  Triggered when live GPU storage
        # (allocated - freed) grows by the threshold since the last GC,
        # indicating biased-refcount tensors are accumulating.
        _is_ft_build = hasattr(sys, "_is_gil_enabled")
        self._gc_on_delete = _is_ft_build
        self._gc_live_bytes_at_last_gc: int = 0
        self._gc_idle_handle = None

    async def start(self, loop: asyncio.AbstractEventLoop) -> None:
        """
        Start the bidirectional stream.

        Args:
            loop: Event loop to run stream tasks on
        """
        if self._started:
            return

        self._loop = loop
        self._request_queue = asyncio.Queue()
        self._shutdown_event = asyncio.Event()

        # Start sender and receiver tasks
        self._sender_task = asyncio.create_task(self._sender_loop())
        self._receiver_task = asyncio.create_task(self._receiver_loop())
        self._started = True

    def _enqueue(self, pending, stream_request):
        """Enqueue a single request. Must run on the event loop thread."""
        if pending is not None:
            self._pending.append(pending)
        self._request_queue.put_nowait(stream_request)

    def _enqueue_batch(self, pending, stream_requests):
        """Enqueue multiple requests with one pending entry. Must run on the event loop thread."""
        for request in stream_requests:
            self._request_queue.put_nowait(request)
        if pending is not None:
            self._pending.append(pending)

    # Flush threshold: when this many ops are buffered, flush immediately.
    # Override with SKYTORCH_BATCH_THRESHOLD env var.
    _BATCH_FLUSH_THRESHOLD = int(os.environ.get("SKYTORCH_BATCH_THRESHOLD", "64"))

    # Coalescing delay: defer partial-batch flush by this many seconds,
    # allowing more ops to accumulate and amortize per-message gRPC overhead.
    # At C++ dispatch rate (~17 us/op), 1 ms ≈ 59 ops per batch.
    # Sync points force an immediate flush, so this never adds sync latency.
    # Override with SKYTORCH_BATCH_COALESCE_MS env var (in milliseconds).
    _BATCH_COALESCE_DELAY = float(os.environ.get("SKYTORCH_BATCH_COALESCE_MS", "2")) / 1000.0

    # Delete coalescing: defer delete flushes by this many seconds, allowing
    # deletes from multiple GC waves to accumulate into a single message.
    # Override with SKYTORCH_DELETE_COALESCE_MS env var (in milliseconds).
    _DELETE_COALESCE_DELAY = float(os.environ.get("SKYTORCH_DELETE_COALESCE_MS", "10")) / 1000.0

    # Safety threshold: flush deletes immediately when this many IDs accumulate,
    # preventing unbounded GPU memory retention during large GC events.
    # Override with SKYTORCH_DELETE_THRESHOLD env var.
    _DELETE_FLUSH_THRESHOLD = int(os.environ.get("SKYTORCH_DELETE_THRESHOLD", "256"))

    # GC reconciliation threshold: run gc.collect(0) when live GPU storage
    # (allocated - freed) grows by this many bytes since the last GC.
    # This directly measures biased-refcount tensor accumulation — tensors
    # allocated but not freed because deferred decrefs haven't been reconciled.
    # Default 1 GiB.  Only active on the free-threaded build.
    # Override with SKYTORCH_GC_DELETE_BYTES env var.
    _GC_DELETE_BYTES_THRESHOLD = int(os.environ.get("SKYTORCH_GC_DELETE_BYTES", str(1024 ** 3)))

    def _enqueue_execute_aten(self, request: service_pb2.ExecuteAtenRequest) -> None:
        """Buffer a protobuf execute_aten request for batching. Must run on the event loop thread."""
        self._batch_buffer.append(request)
        if len(self._batch_buffer) >= self._BATCH_FLUSH_THRESHOLD:
            self._flush_batch()
        elif not self._flush_scheduled:
            self._flush_scheduled = True
            self._batch_flush_timer = self._loop.call_later(
                self._BATCH_COALESCE_DELAY, self._flush_batch
            )

    def _enqueue_execute_aten_bytes(self, raw_bytes: bytes) -> None:
        """Buffer a raw binary execute_aten request for batching. Must run on the event loop thread."""
        self._raw_batch_buffer.append(raw_bytes)
        if len(self._raw_batch_buffer) >= self._BATCH_FLUSH_THRESHOLD:
            self._flush_raw_batch()
        elif not self._raw_flush_scheduled:
            self._raw_flush_scheduled = True
            self._raw_flush_timer = self._loop.call_later(
                self._BATCH_COALESCE_DELAY, self._flush_raw_batch
            )

    def _flush_batch(self) -> None:
        """Flush buffered execute_aten requests as a batch. Must run on the event loop thread."""
        self._flush_scheduled = False
        if self._batch_flush_timer is not None:
            self._batch_flush_timer.cancel()
            self._batch_flush_timer = None
        if not self._batch_buffer:
            return

        if len(self._batch_buffer) == 1:
            # Single request: send as regular execute_aten to avoid wrapping overhead
            stream_request = service_pb2.StreamRequest(execute_aten=self._batch_buffer[0])
        else:
            # Multiple requests: send as batch
            batch = service_pb2.BatchedExecuteAtenRequest(operations=self._batch_buffer)
            stream_request = service_pb2.StreamRequest(batched_execute_aten=batch)
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug(f"Flushing batch of {len(self._batch_buffer)} execute_aten ops")

        self._batch_buffer = []
        self._request_queue.put_nowait(stream_request)

    def _flush_raw_batch(self, piggyback_deletes: bool = True) -> None:
        """Flush buffered raw binary execute_aten requests. Must run on the event loop thread."""
        self._raw_flush_scheduled = False
        if self._raw_flush_timer is not None:
            self._raw_flush_timer.cancel()
            self._raw_flush_timer = None
        if not self._raw_batch_buffer:
            return

        if PROFILING_ENABLED:
            from skytorch.torch.profiler import ClientProfiler

            _prof = ClientProfiler.get()
            _batch_size = len(self._raw_batch_buffer)
            _prof.batch_count.add_count()
            _prof.batch_size_total += _batch_size
            if _batch_size > _prof.batch_size_max:
                _prof.batch_size_max = _batch_size

        if len(self._raw_batch_buffer) == 1:
            # Single request: send as raw_execute_aten
            stream_request = service_pb2.StreamRequest(raw_execute_aten=self._raw_batch_buffer[0])
        else:
            # Multiple requests: concatenate into raw_batched_execute_aten
            # Each request is prefixed with its length as uint32 LE
            import struct

            parts = []
            for raw in self._raw_batch_buffer:
                parts.append(struct.pack("<I", len(raw)))
                parts.append(raw)
            stream_request = service_pb2.StreamRequest(raw_batched_execute_aten=b"".join(parts))
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug(
                    f"Flushing raw batch of {len(self._raw_batch_buffer)} execute_aten ops"
                )

        self._raw_batch_buffer = []
        self._request_queue.put_nowait(stream_request)

        # Piggyback deferred deletes: the batch (which may reference these
        # tensors) was just sent, so it's safe to send deletes now.  This
        # avoids a separate flush that would fragment future batches.
        # Suppressed during sync flushes (get_scalar/get_tensor) where
        # deletes must go AFTER the sync request to avoid races.
        if piggyback_deletes:
            self._send_deferred_deletes()

    async def _sender_loop(self) -> None:
        """Send requests from queue to the stream."""
        try:

            async def request_generator():
                while not self._closing:
                    request = await self._request_queue.get()
                    if request is None:
                        break
                    yield request
                    # Drain ready items without re-awaiting
                    while not self._request_queue.empty():
                        request = self._request_queue.get_nowait()
                        if request is None:
                            return
                        yield request

            # Start the bidirectional stream
            self._stream_call = self._stub.StreamOperations(
                request_generator(), metadata=self._metadata
            )

            # Wait for shutdown signal
            await self._shutdown_event.wait()

        except grpc.RpcError as e:
            logger.error(f"Stream sender error: {e}")
            self._fail_all_pending(str(e))
        except Exception as e:
            logger.error(f"Unexpected stream sender error: {e}")
            self._fail_all_pending(str(e))

    async def _receiver_loop(self) -> None:
        """Receive responses from the stream and resolve futures in order."""
        try:
            # Wait for stream to be established
            while self._stream_call is None and not self._closing:
                await asyncio.sleep(0.01)

            if self._stream_call is None:
                return

            async for response in self._stream_call:
                if not self._pending:
                    logger.warning("Received response but no pending requests")
                    continue
                pending = self._pending.popleft()

                # Check for deferred error from fire-and-forget operations
                if response.error_message and not response.success:
                    error = RuntimeError(f"Stream operation failed: {response.error_message}")
                    if pending.future is not None:
                        pending.future.set_exception(error)
                    self.record_error(error)
                elif pending.future is not None:
                    pending.future.set_result(response)

        except grpc.RpcError as e:
            if not self._closing:
                logger.error(f"Stream receiver error: {e}")
                self._fail_all_pending(str(e))
        except Exception as e:
            if not self._closing:
                logger.error(f"Unexpected stream receiver error: {e}")
                self._fail_all_pending(str(e))

    def record_error(self, error: Exception) -> None:
        """Record error for deferred raising. First-error-wins semantics."""
        if self._last_error is None:
            self._last_error = error

    def check_error(self) -> None:
        """Raise deferred error if any, then clear it."""
        error = self._last_error
        if error is not None:
            self._last_error = None
            raise error

    def _fail_all_pending(self, error_message: str) -> None:
        """Fail all pending requests with an error."""
        for pending in self._pending:
            if pending.future is not None and not pending.future.done():
                pending.future.set_exception(RuntimeError(f"Stream error: {error_message}"))
        self._pending.clear()

    def _drain_cpp_into_mt_ops(self) -> None:
        """Drain C++ raw submit buffer into _mt_ops. Caller must hold _mt_lock.

        Uses C++ flush_raw_batches to build batched payloads in C++,
        reducing event loop GIL-holding time from ~200-500us (Python
        drain + struct.pack + bytes.join) to ~5-10us (C++ concatenation).
        """
        if not self._cpp_submit_available:
            return
        try:
            from skytorch.torch.backend._C import _flush_raw_batches

            batches = _flush_raw_batches(self._BATCH_FLUSH_THRESHOLD)
            for batch_bytes, count in batches:
                self._mt_ops.append(("cpp_batch", (batch_bytes, count)))
        except (ImportError, AttributeError):
            pass

    def _drain_mt_ops(self) -> None:
        """Drain the main-thread ops buffer. Must run on the event loop thread."""
        with self._mt_lock:
            self._drain_cpp_into_mt_ops()
            if not self._mt_ops:
                self._mt_wake_pending = False
                return
            batch = self._mt_ops
            self._mt_ops = []
            self._mt_wake_pending = False
        self._process_mt_ops(batch)

        # Schedule an idle-GC: if no new ops arrive within 100ms, run
        # gc.collect(0) to reconcile any remaining biased refcounts.
        # This handles the gap between the last token of a turn and the
        # next user prompt, where no allocations trigger automatic GC.
        # Cancelled and rescheduled on every _drain_mt_ops call, so it
        # only fires during genuine idle periods.
        if self._gc_on_delete:
            if self._gc_idle_handle is not None:
                self._gc_idle_handle.cancel()
            self._gc_idle_handle = self._loop.call_later(0.1, self._run_idle_gc)

    def _run_idle_gc(self) -> None:
        """Reconcile biased refcounts during idle periods."""
        self._gc_idle_handle = None
        gc.collect(0)

    def _setup_cpp_submit(self) -> None:
        """Set up the C++ raw submit buffer for bypassing Python on the fast path.

        Must be called after the stream is started (self._loop is set).
        The C++ buffer is used only by dispatch_cached_aten's fast path
        (no new tensors). All Python-side submit_execute_aten_bytes calls
        continue through _mt_ops for proven thread-safety.
        """
        if not self._loop:
            return
        try:
            from skytorch.torch.backend._C import _setup_cpp_submit

            _setup_cpp_submit(self._loop.call_soon_threadsafe, self._drain_cpp_raw)
            self._cpp_submit_available = True
        except (ImportError, AttributeError):
            pass

    def _drain_cpp_raw(self) -> None:
        """Drain the C++ raw submit buffer. Must run on the event loop thread.

        Delegates to _drain_mt_ops so that C++ buffer ops are routed through
        _mt_ops, preserving FIFO ordering with "req" entries (update_tensor,
        delete_tensors, etc.). Without this, a C++ op submitted after a "req"
        entry could be drained directly into _raw_batch_buffer, then flushed
        ahead of the "req" by _enqueue_with_flush.
        """
        self._drain_mt_ops()

    def _process_mt_ops(self, ops: list[tuple]) -> None:
        """Process a list of (type, data) ops in order. Must run on the event loop thread."""
        for op_type, data in ops:
            if op_type == "raw":
                self._enqueue_execute_aten_bytes(data)
            elif op_type == "cpp_batch":
                # Pre-batched by C++ flush_raw_batches: (bytes, count).
                # Flush any pending _raw_batch_buffer first to preserve FIFO
                # ordering with prior "raw" entries.
                self._flush_raw_batch()
                batch_bytes, count = data
                if PROFILING_ENABLED:
                    from skytorch.torch.profiler import ClientProfiler

                    _prof = ClientProfiler.get()
                    _prof.batch_count.add_count()
                    _prof.batch_size_total += count
                    if count > _prof.batch_size_max:
                        _prof.batch_size_max = count
                if count == 1:
                    req = service_pb2.StreamRequest(raw_execute_aten=batch_bytes)
                else:
                    req = service_pb2.StreamRequest(raw_batched_execute_aten=batch_bytes)
                self._request_queue.put_nowait(req)
            elif op_type == "aten":
                self._enqueue_execute_aten(data)
            elif op_type == "req":
                self._enqueue_with_flush(data)
            elif op_type == "ff":
                self._deferred_delete_ids.extend(data)
        # Deferred deletes piggyback on the next natural batch flush:
        # _flush_raw_batch calls _send_deferred_deletes after sending the
        # batch, so deletes ride with ATen ops without fragmenting batches.
        # This timer is a safety net for when no ATen ops are pending
        # (e.g., the turn just ended and only deletes remain).
        if self._deferred_delete_ids and not self._deferred_flush_scheduled:
            self._deferred_flush_scheduled = True
            self._deferred_flush_timer = self._loop.call_later(
                self._BATCH_COALESCE_DELAY, self._flush_deferred
            )

    def _flush_mt_ops(self) -> None:
        """Drain main-thread ops buffer at sync points. Must run on the event loop thread."""
        with self._mt_lock:
            # Drain C++ raw submit buffer INTO _mt_ops (not _raw_batch_buffer)
            # to preserve FIFO ordering with "req" entries like update_tensor.
            # Using _drain_cpp_raw() here would put C++ ops into _raw_batch_buffer
            # before _mt_ops is processed, causing _enqueue_with_flush to flush
            # those ops ahead of the "req" entries that should precede them.
            self._drain_cpp_into_mt_ops()
            if not self._mt_ops:
                self._mt_wake_pending = False
                return
            batch = self._mt_ops
            self._mt_ops = []
            self._mt_wake_pending = False
        self._process_mt_ops(batch)

    def _send_deferred_deletes(self) -> None:
        """Send accumulated deferred delete IDs without flushing batch buffers.

        Caller must ensure that any ATen ops referencing the deleted tensors
        have already been sent (either via batch flush or because they were
        enqueued before this call).

        On the free-threaded build, periodically runs gc.collect(0) before
        sending to reconcile biased refcounts from non-owning threads
        (e.g., AsyncStreamer worker).  This is driven by actual delete
        volume, not arbitrary time intervals.

        Must run on the event loop thread.
        """
        self._deferred_flush_scheduled = False
        if self._deferred_flush_timer is not None:
            self._deferred_flush_timer.cancel()
            self._deferred_flush_timer = None
        if not self._deferred_delete_ids:
            return

        # On the free-threaded build, reconcile biased refcounts when live
        # GPU storage grows significantly — indicating biased-refcount
        # tensors are accumulating (allocated but not freed because their
        # deferred decrefs haven't been reconciled).
        if self._gc_on_delete:
            from skytorch.torch.backend._storage import storage_manager

            live = storage_manager.live_storage_bytes
            if live - self._gc_live_bytes_at_last_gc >= self._GC_DELETE_BYTES_THRESHOLD:
                gc.collect(0)
                self._gc_live_bytes_at_last_gc = storage_manager.live_storage_bytes

        ids = self._deferred_delete_ids
        self._deferred_delete_ids = []
        request = service_pb2.DeleteTensorsRequest(tensor_ids=ids)
        stream_request = service_pb2.StreamRequest(delete_tensors=request)
        self._request_queue.put_nowait(stream_request)
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Flushed {len(ids)} deferred deletes")

    def _flush_deferred(self) -> None:
        """Flush deferred delete tensor IDs, ensuring batch buffers go first.

        Flushes any pending ATen ops before sending deletes so that ops
        referencing these tensors reach the server first (FIFO ordering).
        This includes draining the C++ raw submit buffer — ops submitted
        before the tensor was freed may still be in the buffer when the
        deferred delete timer fires.

        Must run on the event loop thread.
        """
        # Drain C++ raw submit buffer directly into the request queue.
        # This avoids _drain_mt_ops (which causes re-entrancy through
        # _process_mt_ops → _enqueue_with_flush → _flush_deferred).
        # FIFO ordering between C++ and Python ops is not critical here —
        # the ops referencing deleted tensors are in the C++ buffer, and
        # they must reach the server before the delete message.
        if self._cpp_submit_available:
            try:
                from skytorch.torch.backend._C import _flush_raw_batches

                with self._mt_lock:
                    batches = _flush_raw_batches(self._BATCH_FLUSH_THRESHOLD)
                for batch_bytes, count in batches:
                    if count == 1:
                        req = service_pb2.StreamRequest(raw_execute_aten=batch_bytes)
                    else:
                        req = service_pb2.StreamRequest(
                            raw_batched_execute_aten=batch_bytes
                        )
                    self._request_queue.put_nowait(req)
            except (ImportError, AttributeError):
                pass
        # Flush remaining batch buffers (Python fallback path ops).
        self._flush_batch()
        self._flush_raw_batch()
        self._send_deferred_deletes()

    def defer_delete_ids(self, tensor_ids: list[int]) -> None:
        """Append tensor IDs for deferred batch deletion. Must run on event loop thread.

        Called directly from the GC callback (via call_soon_threadsafe) to bypass
        the ensure_future → coroutine → _mt_ops roundtrip. IDs accumulate until
        the coalesce timer fires, the threshold is hit, or a sync point flushes.
        """
        self._deferred_delete_ids.extend(tensor_ids)
        if len(self._deferred_delete_ids) >= self._DELETE_FLUSH_THRESHOLD:
            self._flush_deferred()
        elif not self._deferred_flush_scheduled:
            self._deferred_flush_scheduled = True
            self._deferred_flush_timer = self._loop.call_later(
                self._DELETE_COALESCE_DELAY, self._flush_deferred
            )

    def _enqueue_with_flush(self, stream_request):
        """Flush any pending batch, then enqueue a request. Must run on the event loop thread."""
        self._flush_batch()
        self._flush_raw_batch()
        self._flush_deferred()
        self._request_queue.put_nowait(stream_request)

    def _submit_request(self, stream_request: service_pb2.StreamRequest, request_type: str) -> None:
        """
        Submit a fire-and-forget request to the stream (callable from any thread).

        Routes through the unified main-thread ops buffer to preserve FIFO
        ordering with submit_execute_aten_bytes calls.

        Args:
            stream_request: StreamRequest to submit
            request_type: Type of request for tracking
        """
        with self._mt_lock:
            self._drain_cpp_into_mt_ops()
            self._mt_ops.append(("req", stream_request))
            if not self._mt_wake_pending:
                self._mt_wake_pending = True
                self._loop.call_soon_threadsafe(self._drain_mt_ops)

    def submit_execute_aten(self, request: service_pb2.ExecuteAtenRequest) -> None:
        """Submit a fire-and-forget execute_aten request (callable from any thread).

        Requests are buffered and flushed as a batch when the event loop
        processes pending callbacks, reducing per-operation gRPC overhead.
        """
        if logger.isEnabledFor(logging.DEBUG):
            input_tensor_ids = [
                arg.tensor.tensor_id for arg in request.args if arg.WhichOneof("value") == "tensor"
            ]
            output_tensor_ids = [ref.tensor_id for ref in request.outputs]
            logger.debug(
                f"Executing {request.op_name} | "
                f"inputs={input_tensor_ids} | outputs={output_tensor_ids}"
            )
        with self._mt_lock:
            self._drain_cpp_into_mt_ops()
            self._mt_ops.append(("aten", request))
            if not self._mt_wake_pending:
                self._mt_wake_pending = True
                self._loop.call_soon_threadsafe(self._drain_mt_ops)

    def submit_execute_aten_bytes(self, raw_bytes: bytes) -> None:
        """Submit a fire-and-forget binary-serialized execute_aten request.

        Like submit_execute_aten but takes pre-serialized bytes from
        the C++ request builder, avoiding Python protobuf overhead entirely.
        Callable from any thread.

        Ops are collected in a unified main-thread buffer along with other
        submit calls, preserving FIFO ordering while reducing
        call_soon_threadsafe wake-ups by ~98%.
        """
        with self._mt_lock:
            self._drain_cpp_into_mt_ops()
            self._mt_ops.append(("raw", raw_bytes))
            if not self._mt_wake_pending:
                self._mt_wake_pending = True
                self._loop.call_soon_threadsafe(self._drain_mt_ops)

    def submit_copy_tensor(self, request: service_pb2.CopyTensorRequest) -> None:
        """Submit a fire-and-forget copy_tensor request (callable from any thread).

        Serialized as binary with a 0xFE marker byte prefix:
        [0xFE][src_id:u64][dst_id:u64][has_src:u8][src_meta...][has_dst:u8][dst_meta...]

        Where metadata (if present) uses the same binary format as ATen op
        tensor metadata, enabling GIL-free server-side parsing.
        """
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                f"Copying tensor {request.src_tensor_id} " f"to tensor {request.dst_tensor_id}"
            )
        parts = [
            struct.pack("<BQQ", 0xFE, request.src_tensor_id, request.dst_tensor_id),
        ]
        for field in ("src_metadata", "dst_metadata"):
            if request.HasField(field):
                meta = getattr(request, field)
                parts.append(b"\x01")
                parts.append(_serialize_tensor_metadata_binary(meta))
            else:
                parts.append(b"\x00")
        self.submit_execute_aten_bytes(b"".join(parts))

    def submit_register_tensors(self, request: service_pb2.RegisterTensorsRequest) -> None:
        """Submit a fire-and-forget register_tensors request (callable from any thread)."""
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Registering {len(request.registrations)} tensors")
        stream_request = service_pb2.StreamRequest(register_tensors=request)
        self._submit_request(stream_request, "register_tensors")

    def submit_delete_tensors(self, request: service_pb2.DeleteTensorsRequest) -> None:
        """Submit a fire-and-forget delete_tensors request (callable from any thread).

        Delete IDs are routed through the deferred "ff" path, accumulating in
        _deferred_delete_ids and batched into a single DeleteTensorsRequest on
        the next flush (batch timer, threshold, or sync point), reducing
        per-message gRPC overhead.
        """
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Deleting tensors {list(request.tensor_ids)}")
        tensor_ids = list(request.tensor_ids)
        with self._mt_lock:
            self._drain_cpp_into_mt_ops()
            self._mt_ops.append(("ff", tensor_ids))
            if not self._mt_wake_pending:
                self._mt_wake_pending = True
                self._loop.call_soon_threadsafe(self._drain_mt_ops)

    def submit_update_tensor(self, request: service_pb2.UpdateTensorRequest) -> None:
        """Submit a fire-and-forget update_tensor request (callable from any thread)."""
        data_size = len(request.data)

        if logger.isEnabledFor(logging.DEBUG):
            chunked = data_size > STREAM_CHUNK_SIZE
            msg = (
                f"Updating tensor {request.tensor_id} | "
                f"shape={list(request.shape)} | "
                f"size={data_size} bytes"
            )
            if chunked:
                num_chunks = (data_size + STREAM_CHUNK_SIZE - 1) // STREAM_CHUNK_SIZE
                msg += f" | chunks={num_chunks}"
            logger.debug(msg)

        if data_size <= STREAM_CHUNK_SIZE:
            stream_request = service_pb2.StreamRequest(update_tensor=request)
            self._submit_request(stream_request, "update_tensor")
        else:
            self._submit_chunked_update_tensor(request)

    def _submit_chunked_update_tensor(self, request: service_pb2.UpdateTensorRequest) -> None:
        """Submit a large update_tensor request as multiple fire-and-forget chunks."""
        data = request.data
        total_size = len(data)
        total_chunks = (total_size + STREAM_CHUNK_SIZE - 1) // STREAM_CHUNK_SIZE

        mv = memoryview(data)
        stream_requests = []

        for chunk_number in range(total_chunks):
            start = chunk_number * STREAM_CHUNK_SIZE
            end = min(start + STREAM_CHUNK_SIZE, total_size)
            chunk_data = bytes(mv[start:end])

            chunk_request = service_pb2.UpdateTensorRequest(
                tensor_id=request.tensor_id,
                data=chunk_data,
            )
            if chunk_number == 0:
                chunk_request.shape.extend(request.shape)
                chunk_request.dtype = request.dtype
                chunk_request.stride.extend(request.stride)
                chunk_request.storage_offset = request.storage_offset
                if request.HasField("metadata"):
                    chunk_request.metadata.CopyFrom(request.metadata)

            stream_requests.append(
                service_pb2.StreamRequest(
                    update_tensor=chunk_request,
                    chunk_number=chunk_number,
                    total_chunks=total_chunks,
                    total_bytes=total_size if chunk_number == 0 else 0,
                )
            )

        # Route chunks through _mt_ops to preserve FIFO ordering with
        # execute_aten ops. Previously, chunks were enqueued directly on
        # _request_queue via call_soon_threadsafe, but the callback flushed
        # _mt_ops first — sending execute_aten ops that reference the tensor
        # BEFORE the update chunks that create it on the server.
        with self._mt_lock:
            self._drain_cpp_into_mt_ops()
            for request in stream_requests:
                self._mt_ops.append(("req", request))
            if not self._mt_wake_pending:
                self._mt_wake_pending = True
                self._loop.call_soon_threadsafe(self._drain_mt_ops)

    def submit_get_tensor(self, request: service_pb2.GetTensorRequest) -> asyncio.Future:
        """
        Submit a get_tensor request.

        Must be called from the event loop thread.

        Args:
            request: GetTensorRequest to submit

        Returns:
            Awaitable that resolves to StreamResponse with tensor data,
            after checking for deferred errors from prior operations
        """
        if self._loop is None:
            raise RuntimeError("StreamManager not started")

        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Getting tensor {request.tensor_id}")

        if PROFILING_ENABLED:
            _t_flush_start = time.perf_counter_ns()
            _mt_ops_count = len(self._mt_ops)

        # Flush any pending batches before sync operation.
        # Suppress deferred delete piggyback — deletes must go AFTER
        # the sync request to avoid "Tensor not found" races.
        self._flush_mt_ops()
        self._flush_batch()
        self._flush_raw_batch(piggyback_deletes=False)

        if PROFILING_ENABLED:
            from skytorch.torch.profiler import ClientProfiler

            _t_flush_end = time.perf_counter_ns()
            _prof = ClientProfiler.get()
            _prof.sync_flush.add(_t_flush_end - _t_flush_start)
            _prof.sync_mt_ops_total += _mt_ops_count
            _qdepth = self._request_queue.qsize()
            _prof.sync_queue_depth_total += _qdepth
            if _qdepth > _prof.sync_queue_depth_max:
                _prof.sync_queue_depth_max = _qdepth

        future = self._loop.create_future()
        stream_request = service_pb2.StreamRequest(get_tensor=request)

        self._pending.append(
            PendingRequest(
                future=future,
                request_type="get_tensor",
            )
        )
        self._request_queue.put_nowait(stream_request)

        # Flush deferred deletes AFTER the sync request so the server
        # processes the query before any deletions.
        self._flush_deferred()

        if PROFILING_ENABLED:
            _t_enqueue = time.perf_counter_ns()

        async def _with_error_check():
            try:
                response = await future
            except RuntimeError:
                # If there's a prior deferred error (e.g., OOM from a
                # fire-and-forget op), it's the root cause — raise it
                # instead of the symptom (e.g., "Tensor not found").
                # A speculative get_scalar may have consumed the server's
                # deferred error, leaving subsequent syncs with a
                # confusing secondary error.
                self.check_error()
                raise
            if PROFILING_ENABLED:
                _t_resolved = time.perf_counter_ns()
                _wait_ns = _t_resolved - _t_enqueue
                _prof.sync_wait.add(_wait_ns)
                # Decompose sync_wait using server-provided timing
                _server_total = response.server_backlog_ns + response.server_handle_ns
                _network_rtt = _wait_ns - _server_total
                _prof.sync_network_rtt.add(max(0, _network_rtt))
                _prof.sync_server_backlog.add(response.server_backlog_ns)
                _prof.sync_server_handle.add(response.server_handle_ns)
            # FIFO ordering guarantees all prior fire-and-forget errors are recorded
            self.check_error()
            return response

        return _with_error_check()

    def submit_get_scalar(self, request: service_pb2.GetScalarRequest) -> asyncio.Future:
        """
        Submit a get_scalar request.

        Must be called from the event loop thread.

        Args:
            request: GetScalarRequest to submit

        Returns:
            Awaitable that resolves to StreamResponse with scalar value,
            after checking for deferred errors from prior operations
        """
        if self._loop is None:
            raise RuntimeError("StreamManager not started")

        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Getting scalar for tensor {request.tensor_id}")

        if PROFILING_ENABLED:
            _t_flush_start = time.perf_counter_ns()
            _mt_ops_count = len(self._mt_ops)

        # Flush any pending batches before sync operation.
        # Suppress deferred delete piggyback — deletes must go AFTER
        # the sync request to avoid "Tensor not found" races.
        self._flush_mt_ops()
        self._flush_batch()
        self._flush_raw_batch(piggyback_deletes=False)

        if PROFILING_ENABLED:
            from skytorch.torch.profiler import ClientProfiler

            _t_flush_end = time.perf_counter_ns()
            _prof = ClientProfiler.get()
            _prof.sync_flush.add(_t_flush_end - _t_flush_start)
            _prof.sync_mt_ops_total += _mt_ops_count
            _qdepth = self._request_queue.qsize()
            _prof.sync_queue_depth_total += _qdepth
            if _qdepth > _prof.sync_queue_depth_max:
                _prof.sync_queue_depth_max = _qdepth

        future = self._loop.create_future()
        stream_request = service_pb2.StreamRequest(get_scalar=request)

        self._pending.append(
            PendingRequest(
                future=future,
                request_type="get_scalar",
            )
        )
        self._request_queue.put_nowait(stream_request)

        # Flush deferred deletes AFTER the sync request so the server
        # processes the query before any deletions.
        self._flush_deferred()

        if PROFILING_ENABLED:
            _t_enqueue = time.perf_counter_ns()

        async def _with_error_check():
            try:
                response = await future
            except RuntimeError:
                # If there's a prior deferred error (e.g., OOM from a
                # fire-and-forget op), it's the root cause — raise it
                # instead of the symptom (e.g., "Tensor not found").
                # A speculative get_scalar may have consumed the server's
                # deferred error, leaving subsequent syncs with a
                # confusing secondary error.
                self.check_error()
                raise
            if PROFILING_ENABLED:
                _t_resolved = time.perf_counter_ns()
                _wait_ns = _t_resolved - _t_enqueue
                _prof.sync_wait.add(_wait_ns)
                # Decompose sync_wait using server-provided timing
                _server_total = response.server_backlog_ns + response.server_handle_ns
                _network_rtt = _wait_ns - _server_total
                _prof.sync_network_rtt.add(max(0, _network_rtt))
                _prof.sync_server_backlog.add(response.server_backlog_ns)
                _prof.sync_server_handle.add(response.server_handle_ns)
            # FIFO ordering guarantees all prior fire-and-forget errors are recorded
            self.check_error()
            return response

        return _with_error_check()

    def submit_execute_module_forward(
        self, request: service_pb2.ExecuteModuleForwardRequest
    ) -> asyncio.Future:
        """
        Submit an execute_module_forward request (sync).

        Must be called from the event loop thread. Flushes pending batches
        before sending to ensure all prior ops complete first.

        Args:
            request: ExecuteModuleForwardRequest to submit

        Returns:
            Awaitable that resolves to StreamResponse with module forward result
        """
        if self._loop is None:
            raise RuntimeError("StreamManager not started")

        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                f"Executing module forward: model_id={request.model_id} "
                f"module={request.module_path} inputs={list(request.input_tensor_ids)}"
            )

        # Flush any pending batches before sync operation.
        # Suppress deferred delete piggyback — deletes must go AFTER
        # the sync request to avoid "Tensor not found" races.
        self._flush_mt_ops()
        self._flush_batch()
        self._flush_raw_batch(piggyback_deletes=False)

        future = self._loop.create_future()
        stream_request = service_pb2.StreamRequest(execute_module_forward=request)

        self._pending.append(
            PendingRequest(
                future=future,
                request_type="execute_module_forward",
            )
        )
        self._request_queue.put_nowait(stream_request)

        # Flush deferred deletes AFTER the sync request.
        self._flush_deferred()

        async def _with_error_check():
            response = await future
            self.check_error()
            return response

        return _with_error_check()

    def submit_execute_module_forward_ff(
        self, request: service_pb2.ExecuteModuleForwardRequest
    ) -> None:
        """Submit a fire-and-forget execute_module_forward request (callable from any thread).

        Used when output tensors are pre-allocated by the client (cached shapes).
        The request must have output_tensor_ids and output_metadata populated.

        Serialized as binary with a 0xFF marker byte prefix:
        [0xFF][model_id:u64][path_len:u16][path:utf8][n_in:u8][in_ids:u64...][n_out:u8][out_ids:u64...]

        Routed through the raw binary ATen batching pipeline, so module forward
        ops are batched with surrounding ATen ops.
        """
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                f"Executing module forward (ff): model_id={request.model_id} "
                f"module={request.module_path} outputs={list(request.output_tensor_ids)}"
            )
        path_bytes = request.module_path.encode("utf-8")
        n_in = len(request.input_tensor_ids)
        n_out = len(request.output_tensor_ids)
        raw_bytes = struct.pack(
            f"<BQH{len(path_bytes)}sB{n_in}QB{n_out}Q",
            0xFF,
            request.model_id,
            len(path_bytes),
            path_bytes,
            n_in,
            *request.input_tensor_ids,
            n_out,
            *request.output_tensor_ids,
        )
        self.submit_execute_aten_bytes(raw_bytes)

    def submit_release_model(self, request: service_pb2.ReleaseModelRequest) -> None:
        """Submit a fire-and-forget release_model request (callable from any thread)."""
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Releasing model {request.model_id}")
        stream_request = service_pb2.StreamRequest(release_model=request)
        self._submit_request(stream_request, "release_model")

    async def close(self) -> None:
        """Gracefully close the stream."""
        if not self._started:
            return

        # Flush pending batches before closing
        self._flush_mt_ops()
        self._flush_batch()
        self._flush_raw_batch()
        self._flush_deferred()

        # Clear C++ submit buffer state
        if self._cpp_submit_available:
            self._cpp_submit_available = False
            try:
                from skytorch.torch.backend._C import _clear_cpp_submit

                _clear_cpp_submit()
            except (ImportError, AttributeError):
                pass

        self._closing = True

        # Wake up sender loop
        if self._shutdown_event is not None:
            self._shutdown_event.set()

        # Signal sender to stop — the request_generator yields no more requests,
        # the server sees end-of-stream, runs its finally cleanup, and closes
        # the response stream. The receiver completes when it sees the stream end.
        if self._request_queue is not None:
            self._request_queue.put_nowait(None)

        # Wait for sender to finish (it returns once request_generator ends)
        if self._sender_task is not None:
            try:
                await asyncio.wait_for(self._sender_task, timeout=5.0)
            except (asyncio.CancelledError, asyncio.TimeoutError):
                self._sender_task.cancel()
                try:
                    await self._sender_task
                except asyncio.CancelledError:
                    pass

        # Wait for receiver to finish (server closes stream after cleanup)
        if self._receiver_task is not None:
            try:
                await asyncio.wait_for(self._receiver_task, timeout=5.0)
            except (asyncio.CancelledError, asyncio.TimeoutError):
                if self._stream_call is not None:
                    self._stream_call.cancel()
                self._receiver_task.cancel()
                try:
                    await self._receiver_task
                except asyncio.CancelledError:
                    pass

        # Fail any remaining pending requests
        self._fail_all_pending("Stream closed")

        self._started = False

        if PROFILING_ENABLED:
            from skytorch.torch.profiler import ClientProfiler

            ClientProfiler.get().print_summary()
            ClientProfiler.reset()
