"""B7 unit tests — the async surface + the async in-process transport.

Covers:

* async round-trips vs the core fakes (create → start → stop → delete via
  ``AsyncVmManagerClient`` over ``AsyncInProcessTransport``);
* reads + capabilities + watch_vm on the async surface;
* the executor shuts down on ``aclose`` — no leaked threads (thread-count
  assertion);
* **cancellation across the thread boundary**: cancelling an in-flight
  awaitable signals the ``CancellationToken`` the offloaded op observes → the
  op converges to ``CANCELLED`` / ``OperationCancelled``;
* per-``VmId`` serialization still holds through the offload (the sync engine's
  lock is acquired in the worker thread).
"""

from __future__ import annotations

import asyncio
import threading
from datetime import UTC, datetime
from typing import TYPE_CHECKING

import pytest
from mgf.vm.core.cancellation import CancellationEvent
from mgf.vm.core.contracts import (
    EventGap,
    OperationKind,
    OperationStatus,
    Vm,
    VmEvent,
    VmState,
)
from mgf.vm.core.exceptions import DomainNotFound, OperationCancelled
from mgf.vm.core.fakes.hypervisor import FakeHypervisorBackend
from mgf.vm.core.ids import DomainId, VmId

from mgf.vm.sdk.client import AsyncVmManagerClient
from mgf.vm.sdk.transport.async_inprocess import AsyncInProcessTransport

if TYPE_CHECKING:
    from collections.abc import Iterator

    from mgf.vm.core.cancellation import CancellationToken
    from mgf.vm.core.contracts import VmSpec


def _spec(name: str = "vm-a") -> VmSpec:
    from mgf.vm.core.contracts import VmSpec

    return VmSpec(name=name, cpu=2, memory_bytes=2 << 30)


# -- async round-trip vs fakes ---------------------------------------------


async def test_async_create_start_stop_delete_round_trip() -> None:
    backend = FakeHypervisorBackend()
    transport = AsyncInProcessTransport.from_adapters(backend)
    async with AsyncVmManagerClient(transport) as client:
        vm = await client.create_vm(_spec("rt"))
        assert isinstance(vm, Vm)
        assert vm.spec.name == "rt"
        vm_id = vm.vm_id

        await client.start_vm(vm_id)
        assert (await client.get_vm(vm_id)).status.state is VmState.RUNNING

        await client.stop_vm(vm_id)
        assert (await client.get_vm(vm_id)).status.state is VmState.STOPPED

        await client.delete_vm(vm_id)
        with pytest.raises(DomainNotFound):
            await client.get_vm(vm_id)


async def test_async_clone_backup_restore_round_trip() -> None:
    # Host disk-image ops (D-G18) on the async surface, over the offload pool.
    from mgf.vm.core.contracts import DiskSpec, VmSpec
    from mgf.vm.core.fakes.disk import FakeHostDiskBackend

    backend = FakeHostDiskBackend()
    backend._put_image("/img/vda.qcow2", content=b"root")
    transport = AsyncInProcessTransport.from_adapters(FakeHypervisorBackend(), host_disk=backend)
    spec = VmSpec(
        name="src",
        cpu=1,
        memory_bytes=1 << 30,
        disks=(DiskSpec(size_bytes=1 << 30, source="/img/vda.qcow2"),),
    )
    async with AsyncVmManagerClient(transport) as client:
        src = (await client.create_vm(spec)).vm_id

        clone = await client.clone_vm(src, "clone", dest_dir="/clones")
        assert clone.spec.disks[0].source == "/clones/clone-0.qcow2"

        await client.backup_vm(src, "/b/src.tar")
        restored = await client.restore_vm("/b/src.tar", dest_dir="/restore", new_name="restored")
        assert restored.spec.name == "restored"
        assert restored.spec.disks[0].source == "/restore/disk0.qcow2"


async def test_async_network_management_round_trip() -> None:
    # Network management (D-G15) on the async surface, over the offload pool.
    from mgf.vm.core.contracts import NetworkSpec
    from mgf.vm.core.fakes import FakeNetworkBackend

    hv = FakeHypervisorBackend()
    transport = AsyncInProcessTransport.from_adapters(hv, network=FakeNetworkBackend(hv))
    async with AsyncVmManagerClient(transport) as client:
        profile = await client.define_network(
            NetworkSpec(name="anet", forward_mode="nat", cidr="10.9.0.0/24")
        )
        assert profile.name == "anet"
        await client.start_network("anet")
        nets = {n.name: n for n in await client.list_networks()}
        assert nets["anet"].active is True
        await client.delete_network("anet")
        assert "anet" not in {n.name for n in await client.list_networks()}


async def test_async_reads_and_capabilities() -> None:
    backend = FakeHypervisorBackend()
    transport = AsyncInProcessTransport.from_adapters(backend)
    async with AsyncVmManagerClient(transport) as client:
        await client.create_vm(_spec("one"))
        await client.create_vm(_spec("two"))
        vms = await client.list_vms()
        assert all(isinstance(v, Vm) for v in vms)
        names = sorted(v.spec.name for v in vms)
        assert names == ["one", "two"]
        caps = await client.capabilities()
        assert caps.engine == "fake"
        ops = await client.list_operations()
        assert len(ops) == 2


async def test_async_error_bridge_reraises_typed() -> None:
    backend = FakeHypervisorBackend()
    transport = AsyncInProcessTransport.from_adapters(backend)
    async with AsyncVmManagerClient(transport) as client:
        with pytest.raises(DomainNotFound):
            await client.start_vm(VmId("fake-missing"))


async def test_async_get_operation_round_trips() -> None:
    backend = FakeHypervisorBackend()
    transport = AsyncInProcessTransport.from_adapters(backend)
    async with AsyncVmManagerClient(transport) as client:
        vm = await client.create_vm(_spec("op"))
        ops = await client.list_operations()
        assert len(ops) == 1
        handle = await client.get_operation(ops[0].op_id)
        assert handle.status is OperationStatus.SUCCEEDED
        assert isinstance(vm, Vm)
        assert isinstance(vm.vm_id, VmId)


async def test_async_watch_op_streams_inflight_transitions() -> None:
    """A *concurrent* watcher resolves an in-flight op via the async ``watch_op``.

    The in-process ``submit_op`` runs to terminal in its worker, so its own
    handle is already settled — to exercise the async ``watch_op`` generator we
    subscribe from a second task while the op is still RUNNING (the op_id is
    read from the shared registry).
    """
    backend = _GatedBackend()
    transport = AsyncInProcessTransport.from_adapters(backend, max_workers=4)
    info = transport_create(transport, "inflight")
    vm_id = info.vm_id

    submit_task = asyncio.ensure_future(
        transport.submit_op(OperationKind.VM_START, vm_id, client_op_id="if-1")
    )
    # Wait until the worker is parked inside ``start`` (op registered + RUNNING).
    await asyncio.get_running_loop().run_in_executor(None, backend.entered.wait, 5.0)
    running = [o for o in transport._engine._registry.list() if o.kind is OperationKind.VM_START]
    assert running and running[0].status is OperationStatus.RUNNING
    op_id = running[0].op_id

    seen: list[OperationStatus] = []

    async def watch() -> None:
        async for h in transport.watch_op(op_id):
            seen.append(h.status)

    watch_task = asyncio.ensure_future(watch())
    await asyncio.sleep(0.05)
    backend.release()
    await asyncio.gather(submit_task, watch_task)
    assert OperationStatus.SUCCEEDED in seen
    await transport.aclose()


async def test_async_submit_with_non_event_token_is_honored() -> None:
    """A caller token that isn't a CancellationEvent is threaded as-is.

    The transport can't drive its cancel from the async side, but the op still
    runs and succeeds (the engine honors the token at the backend).
    """

    class _ReadOnlyToken:
        @property
        def cancelled(self) -> bool:
            return False

        def raise_if_cancelled(self) -> None:
            return

        def wait(self, timeout: float | None = None) -> bool:
            return False

    backend = FakeHypervisorBackend()
    transport = AsyncInProcessTransport.from_adapters(backend)
    info = transport_create(transport, "tok")
    vm_id = info.vm_id
    handle = await transport.submit_op(
        OperationKind.VM_START, vm_id, client_op_id="tok-1", token=_ReadOnlyToken()
    )
    assert handle.status is OperationStatus.SUCCEEDED
    await transport.aclose()


class _GatedBackend(FakeHypervisorBackend):
    """``start`` blocks on a gate so the async watch_op bridge sees transitions."""

    def __init__(self) -> None:
        super().__init__()
        self.entered = threading.Event()
        self._gate = threading.Event()

    def release(self) -> None:
        self._gate.set()

    def start(self, d: DomainId, *, token: CancellationToken | None = None) -> None:
        self.entered.set()
        self._gate.wait(timeout=5.0)
        super().start(d, token=token)


# -- async watch_vm --------------------------------------------------------


class _EventingHypervisor(FakeHypervisorBackend):
    def __init__(self, script: list[VmEvent | EventGap]) -> None:
        super().__init__()
        self._script = script

    def watch_events(self, d: DomainId | None = None) -> Iterator[VmEvent | EventGap]:
        return iter(self._script)


async def test_async_watch_vm_yields_events_and_gap() -> None:
    now = datetime.now(UTC)
    script: list[VmEvent | EventGap] = [
        VmEvent(domain_id=DomainId("d-1"), sequence=1, timestamp=now, payload={}),
        EventGap(domain_id=DomainId("d-1"), after_sequence=1, detected_at=now),
        VmEvent(domain_id=DomainId("d-1"), sequence=5, timestamp=now, payload={}),
    ]
    backend = _EventingHypervisor(script)
    transport = AsyncInProcessTransport.from_adapters(backend)
    async with AsyncVmManagerClient(transport) as client:
        out = [ev async for ev in client.watch_vm(VmId("d-1"))]
    assert [type(x).__name__ for x in out] == ["VmEvent", "EventGap", "VmEvent"]


# -- executor lifetime: no leaked threads ----------------------------------


def _async_worker_count() -> int:
    return sum(1 for t in threading.enumerate() if t.name.startswith("mgf-vm-sdk-async"))


async def test_executor_shuts_down_on_aclose_no_leaked_threads() -> None:
    before = _async_worker_count()
    backend = FakeHypervisorBackend()
    transport = AsyncInProcessTransport.from_adapters(backend, max_workers=4)
    client = AsyncVmManagerClient(transport)

    # Drive several ops so worker threads are actually spawned.
    for i in range(6):
        await client.create_vm(_spec(f"vm-{i}"))
    assert _async_worker_count() > before  # pool spun up workers

    await client.aclose()
    # aclose joins the pool (shutdown(wait=True)) — workers are gone.
    assert _async_worker_count() == before
    assert backend.closed
    # Idempotent.
    await client.aclose()


# -- cancellation across the thread boundary -------------------------------


class _BlockingBackend(FakeHypervisorBackend):
    """``start`` blocks until its cancellation token is signalled.

    Models a long-running adapter call: it parks on ``token.wait`` and raises
    ``OperationCancelled`` the instant the token flips — which is exactly what
    the async transport must trigger when the awaiting task is cancelled.
    """

    def __init__(self) -> None:
        super().__init__()
        self.entered = threading.Event()
        self.observed_cancel = False

    def start(self, d: DomainId, *, token: CancellationToken | None = None) -> None:
        self.entered.set()
        assert token is not None, "async transport must thread a token through"
        # Park until cancelled (bounded so a buggy test can't hang forever).
        cancelled = token.wait(timeout=5.0)
        if cancelled:
            self.observed_cancel = True
            token.raise_if_cancelled()  # -> OperationCancelled
        super().start(d, token=token)


async def test_async_cancellation_signals_token_to_worker() -> None:
    backend = _BlockingBackend()
    transport = AsyncInProcessTransport.from_adapters(backend)
    client = AsyncVmManagerClient(transport)
    info = await client.create_vm(_spec("cx"))
    vm_id = info.vm_id

    task = asyncio.ensure_future(client.start_vm(vm_id))
    # Wait until the worker is parked inside the blocking backend call.
    await asyncio.get_running_loop().run_in_executor(None, backend.entered.wait, 5.0)

    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task

    # The worker must have OBSERVED the cancel (token signalled across the
    # boundary), not just been abandoned.
    await asyncio.get_running_loop().run_in_executor(None, _spin_until, backend)
    assert backend.observed_cancel is True
    await client.aclose()


async def test_async_precancelled_token_marks_cancelled() -> None:
    """A pre-cancelled token threaded by the caller → CANCELLED handle."""
    backend = FakeHypervisorBackend()
    transport = AsyncInProcessTransport.from_adapters(backend)
    info = transport_create(transport, "pc")
    vm_id = info.vm_id

    token = CancellationEvent()
    token.cancel()
    handle = await transport.submit_op(
        # VM_START with a pre-cancelled token: the fake raises at entry.
        OperationKind.VM_START,
        vm_id,
        client_op_id="pc-1",
        token=token,
    )
    assert handle.status is OperationStatus.CANCELLED
    async with AsyncVmManagerClient(transport):
        with pytest.raises(OperationCancelled):
            from mgf.vm.sdk.client import VmManagerClient

            VmManagerClient._unpack(handle)


# -- per-VmId serialization still holds through the offload ----------------


async def test_per_vm_serialization_holds_through_offload() -> None:
    backend = _ConcurrencyProbe()
    transport = AsyncInProcessTransport.from_adapters(backend, max_workers=4)
    async with AsyncVmManagerClient(transport) as client:
        info = await client.create_vm(_spec("ser"))
        vm_id = info.vm_id
        await client.start_vm(vm_id)
        # Two ops on the SAME vm, concurrently scheduled, must serialize on the
        # sync engine's per-VM lock (acquired in the worker thread).
        backend.gate = True
        await asyncio.gather(
            client.reboot_vm(vm_id),
            client.reboot_vm(vm_id),
        )
    assert backend.max_active <= 1


# -- helpers ---------------------------------------------------------------


def transport_create(transport: AsyncInProcessTransport, name: str) -> Vm:
    """Create a VM through the wrapped sync engine directly (no event loop)."""
    from mgf.vm.core.contracts import OperationKind
    from mgf.vm.core.contracts import Vm as _Vm

    handle = transport._engine.submit_op(
        OperationKind.VM_CREATE, _spec(name), client_op_id=f"c-{name}"
    )
    result = handle.result
    assert isinstance(result, _Vm)
    return result


def _spin_until(backend: _BlockingBackend) -> None:
    import time

    deadline = time.monotonic() + 5.0
    while time.monotonic() < deadline and not backend.observed_cancel:
        time.sleep(0.01)


class _ConcurrencyProbe(FakeHypervisorBackend):
    """Records peak concurrent ``reboot`` calls to detect lock serialization."""

    def __init__(self) -> None:
        super().__init__()
        self.gate = False
        self._active = 0
        self.max_active = 0
        self._probe_lock = threading.Lock()

    def reboot(
        self,
        d: DomainId,
        *,
        force: bool = False,
        token: CancellationToken | None = None,
    ) -> None:
        if self.gate:
            with self._probe_lock:
                self._active += 1
                self.max_active = max(self.max_active, self._active)
            import time

            time.sleep(0.05)
            with self._probe_lock:
                self._active -= 1
        super().start(d, token=token)
        super().reboot(d, force=force, token=token)
