"""``AsyncInProcessTransport`` — the async mirror of the in-process kernel (B7).

The concrete :class:`~mgf.vm.sdk.transport.base.AsyncTransport` the
:class:`~mgf.vm.sdk.client.AsyncVmManagerClient` runs on for VManager Local.

Why a wrapper, not a rewrite
----------------------------

``mgf.vm.core`` and every adapter (libvirt) are **synchronous** — there is no
async libvirt. The orchestration logic (op registry, per-``VmId`` locking,
state machine, error/cancel bridge) already lives, correct, in
:class:`~mgf.vm.sdk.transport.inprocess.InProcessTransport`. Duplicating it for
async would be two engines to keep in lockstep. So this transport **wraps a
single sync** :class:`InProcessTransport` (it shares that engine — the same
registry, the same per-VM locks) and turns each blocking call into an
``await`` by **offloading it to a thread pool** via
``loop.run_in_executor``. Async here is purely an I/O-concurrency affordance
over a blocking core; the sync engine remains the source of truth.

Executor lifetime
-----------------

The transport **owns** a :class:`concurrent.futures.ThreadPoolExecutor`. It is
created at construction and shut down by :meth:`aclose` (which also closes the
wrapped sync engine). ``aclose`` is idempotent and joins the pool's workers
(``shutdown(wait=True)``), so no threads leak past close (DP-13). The pool is
bounded; per-``VmId`` serialization is unaffected because the *sync engine's*
locks are acquired inside the worker thread — two offloaded ops on the same VM
still serialize on that VM's lock exactly as they do synchronously.

Cancellation across the thread boundary (the flagged hazard)
------------------------------------------------------------

When an awaiting asyncio task is cancelled, the worker thread running the
blocking op does **not** stop on its own — a thread can't be force-killed, and
abandoning it would leak the work (and any backend side effects mid-flight).
Instead, every offloaded mutating call is threaded a
:class:`~mgf.vm.core.cancellation.CancellationToken`. On
``asyncio.CancelledError`` we **signal that token** (``cancel()``) so the
worker observes cancellation at the backend's next yield point and unwinds
cleanly to a ``CANCELLED`` handle — the same D-D7 path the sync surface uses.
We then re-raise ``CancelledError`` to the awaiting task. The caller passes
their own token, or we mint a :class:`~mgf.vm.core.cancellation.CancellationEvent`
to drive the cancel through.
"""

from __future__ import annotations

import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING

from mgf.vm.core.cancellation import CancellationEvent, CancellationToken

from mgf.vm.sdk.transport.inprocess import InProcessTransport

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from mgf.vm.core.contracts import (
        DeviceBackend,
        DomainPolicyBackend,
        EventGap,
        GuestAgentBackend,
        HostDiskBackend,
        HypervisorBackend,
        LifecycleObserver,
        MetricsBackend,
        NetworkBackend,
        OperationHandle,
        OperationKind,
        StorageBackend,
        VmEvent,
    )
    from mgf.vm.core.ids import OperationId

__all__ = ["AsyncInProcessTransport"]

# Sentinel ending the bridged sync-iterator pull (a real item is never this).
_STREAM_DONE = object()


class AsyncInProcessTransport:
    """Async :class:`~mgf.vm.sdk.transport.base.AsyncTransport` over the sync engine.

    Construct via :meth:`from_adapters`. Wraps one
    :class:`~mgf.vm.sdk.transport.inprocess.InProcessTransport` (the shared
    engine) and offloads each blocking call to an owned thread pool. Use as an
    async context manager to guarantee :meth:`aclose`.
    """

    def __init__(
        self,
        engine: InProcessTransport,
        *,
        max_workers: int | None = None,
    ) -> None:
        self._engine = engine
        # thread-name prefix aids leak debugging (the leak-count test reads it).
        self._executor = ThreadPoolExecutor(
            max_workers=max_workers,
            thread_name_prefix="mgf-vm-sdk-async",
        )
        self._closed = False

    @classmethod
    def from_adapters(
        cls,
        hypervisor: HypervisorBackend,
        storage: StorageBackend | None = None,
        device: DeviceBackend | None = None,
        *,
        metrics: MetricsBackend | None = None,
        host_disk: HostDiskBackend | None = None,
        network: NetworkBackend | None = None,
        guest_agent: GuestAgentBackend | None = None,
        policy: DomainPolicyBackend | None = None,
        observers: tuple[LifecycleObserver, ...] = (),
        max_workers: int | None = None,
    ) -> AsyncInProcessTransport:
        """Build an async transport over a fresh sync in-process engine (B7).

        Mirrors :meth:`InProcessTransport.from_adapters`: ``hypervisor`` is
        required, ``storage`` / ``device`` / ``metrics`` / ``host_disk`` /
        ``network`` optional (``metrics`` is the OPTIONAL Wave-B
        :class:`~mgf.vm.core.contracts.MetricsBackend`, D-G13; ``host_disk`` the
        OPTIONAL Wave-R4 :class:`~mgf.vm.core.contracts.HostDiskBackend`, D-G18;
        ``network`` the OPTIONAL Wave-R3
        :class:`~mgf.vm.core.contracts.NetworkBackend`, D-G15), ``observers``
        fire on every VM's state machine. ``max_workers`` bounds the offload
        pool.
        """
        engine = InProcessTransport.from_adapters(
            hypervisor,
            storage,
            device,
            metrics=metrics,
            host_disk=host_disk,
            network=network,
            guest_agent=guest_agent,
            policy=policy,
            observers=observers,
        )
        return cls(engine, max_workers=max_workers)

    # -- offload primitive -------------------------------------------------

    async def _offload(self, func: object) -> object:
        """Run ``func`` (a zero-arg callable) in the pool; await the result."""
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(self._executor, func)  # type: ignore[arg-type]

    # -- AsyncTransport seam: mutations -----------------------------------

    async def submit_op(
        self,
        kind: OperationKind,
        payload: object,
        *,
        client_op_id: str,
        token: object | None = None,
    ) -> OperationHandle:
        """Offload the sync :meth:`InProcessTransport.submit_op` to the pool.

        The sync engine runs the op to terminal inside the worker thread (under
        the VM's lock). On ``asyncio.CancelledError`` — the awaiting task was
        cancelled — we signal the cancellation token so the worker observes it
        and unwinds to ``CANCELLED``, then re-raise. The caller's token is used
        if given; otherwise we mint a :class:`CancellationEvent` purely to carry
        the cancel across the thread boundary.
        """
        signal = _as_event(token)
        effective: object | None = token if token is not None else signal

        def _work() -> OperationHandle:
            return self._engine.submit_op(kind, payload, client_op_id=client_op_id, token=effective)

        future = asyncio.get_running_loop().run_in_executor(self._executor, _work)
        try:
            return await asyncio.shield(future)
        except asyncio.CancelledError:
            # Signal the worker (it can't be force-killed) so it converges to
            # CANCELLED at the backend's next yield point; then propagate.
            if signal is not None:
                signal.cancel()
            future.cancel()
            raise

    async def get_op(self, op_id: OperationId) -> OperationHandle:
        """Offload the sync :meth:`InProcessTransport.get_op` (one-shot poll)."""
        return await self._offload(lambda: self._engine.get_op(op_id))  # type: ignore[return-value]

    async def watch_op(self, op_id: OperationId) -> AsyncIterator[OperationHandle]:
        """Async-iterate ``op_id``'s transitions to terminal (``WatchOp``).

        Bridges the sync :class:`~mgf.vm.sdk.operations.WatchOpStream`: each
        blocking ``next()`` is offloaded to the pool, so the event loop is never
        blocked. The underlying op already ran to terminal in :meth:`submit_op`,
        so this usually yields the terminal handle once and stops.
        """
        sync_iter = await self._offload(lambda: self._engine.watch_op(op_id))
        async for handle in self._bridge_iter(sync_iter):
            yield handle  # type: ignore[misc]

    # -- AsyncTransport seam: reads + streams -----------------------------

    async def read(self, query: object) -> object:
        """Offload the sync :meth:`InProcessTransport.read` to the pool."""
        return await self._offload(lambda: self._engine.read(query))

    async def watch_stream(self, subject: object) -> AsyncIterator[VmEvent | EventGap]:
        """Async-iterate resource events for ``subject`` (``WatchVm``; B6/B7).

        Bridges the sync hypervisor event iterator: each blocking ``next()`` is
        offloaded so the loop stays responsive. :class:`EventGap` items pass
        straight through (not swallowed) so callers reconcile. Closing the async
        generator (``aclose`` / breaking the loop) closes the sync iterator,
        releasing the engine subscription (DP-13).
        """
        sync_iter = await self._offload(lambda: self._engine.watch_stream(subject))
        async for event in self._bridge_iter(sync_iter):
            yield event  # type: ignore[misc]

    async def _bridge_iter(self, sync_iter: object) -> AsyncIterator[object]:
        """Pump a blocking iterator through the pool, one offloaded ``next()`` per item.

        Closes the wrapped iterator on exit (generator ``aclose`` / break) so
        the engine subscription is released deterministically (DP-13).
        """
        iterator = iter(sync_iter)  # type: ignore[call-overload]

        def _next() -> object:
            try:
                return next(iterator)
            except StopIteration:
                return _STREAM_DONE

        try:
            while True:
                item = await self._offload(_next)
                if item is _STREAM_DONE:
                    return
                yield item
        finally:
            close = getattr(iterator, "close", None)
            if callable(close):
                close()

    # -- lifecycle ---------------------------------------------------------

    async def aclose(self) -> None:
        """Shut the offload pool down + close the sync engine. Idempotent (DP-13).

        Joins the pool's workers (``shutdown(wait=True)``) so no threads leak,
        then closes the wrapped engine (releasing its backends).
        """
        if self._closed:
            return
        self._closed = True
        # Offload the (possibly blocking) backend close, then drain the pool.
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(self._executor, self._engine.close)
        await loop.run_in_executor(None, lambda: self._executor.shutdown(wait=True))


def _as_event(token: object | None) -> CancellationEvent | None:
    """Return a :class:`CancellationEvent` we can ``cancel()`` to signal the worker.

    If the caller supplied a token that is itself cancellable (a
    :class:`CancellationEvent`), reuse it. Otherwise mint a fresh event whose
    cancel we drive on ``CancelledError`` — but only when no token was given
    (a caller-supplied non-event token is honored as-is by the engine; we just
    can't drive its cancel from here).
    """
    if token is None:
        return CancellationEvent()
    if isinstance(token, CancellationEvent):
        return token
    return None


# Quiet "imported but unused at runtime" — CancellationToken documents the
# token contract the offload threads through to the backends.
_ = CancellationToken
