"""Device-domain handler — leases + THE ATTACH TRANSACTION (B5).

Two layers live here:

* **Lease lifecycle** (``DEVICE_CLAIM`` / ``DEVICE_RELEASE`` /
  ``DEVICE_TRANSFER``) — thin ``DeviceBackend`` calls returning a ``LeaseId``
  or ``None``. Idempotency + typed errors (``DeviceContended`` /
  ``LeaseNotFound`` / …) are the backend's; they propagate verbatim.

* **The attach transaction** (``contracts.md`` §5) — the cross-backend
  integration the agent normally coordinates, run in-process here:

  ``attach_device(vm_id, device_id)``::

      lease = device.claim(device_id, mode=EXCLUSIVE)   # broker leases it
      try:
          assignment = device.assignment_for(lease)     # → host handle
          hypervisor.attach_device(domain, assignment)  # plug into guest
      except Exception:
          device.release(lease)   # ROLL BACK: broker safe-resets the device
          raise
      # success — the lease stays held; detach releases it later.

  **Rollback contract.** If *anything* after the successful ``claim`` fails
  (``assignment_for`` raising ``LeaseExpired``, the hypervisor raising
  ``DeviceAttachFailed``, …) the lease is released so the broker safe-resets
  the device and we do not leak a lease (a lease leak is a defect,
  ``contracts.md`` §5). The ORIGINAL failure is then re-raised — the
  ``release`` is best-effort cleanup and must not mask the real cause.

  ``detach_device(vm_id, device_id)`` mirrors **in reverse**: hypervisor
  ``detach_device`` first, then ``device.release``. If the hypervisor detach
  fails the lease is **NOT** released — the device is still presented to the
  guest and releasing the lease would ghost it. We re-raise the detach failure
  as a typed :class:`~mgf.vm.core.exceptions.DeviceDetachFailed` ("needs-human"
  — the CP records the discrepancy; ``contracts.md`` §5). Because detach needs
  the lease to release, the transport tracks attach→lease bindings in a small
  in-memory map.

Attach/detach are VM-scoped (they touch a domain) and run under the
transport's per-``VmId`` lock; the bare lease ops are resource-scoped and take
no VM lock.
"""

from __future__ import annotations

import threading
from typing import TYPE_CHECKING, TypeVar

from mgf.vm.core.contracts import ClaimMode, OperationKind
from mgf.vm.core.exceptions import DeviceDetachFailed, TransportError, VmManagerError
from mgf.vm.core.ids import DomainId

from mgf.vm.sdk.queries import (
    ClaimDevicePayload,
    DetachDevicePayload,
    ReleaseDevicePayload,
    TransferLeasePayload,
    VmAttachDevicePayload,
)

if TYPE_CHECKING:
    from mgf.vm.core.cancellation import CancellationToken
    from mgf.vm.core.contracts import DeviceBackend, HypervisorBackend
    from mgf.vm.core.ids import DeviceId, LeaseId, VmId

# Bare lease-lifecycle kinds (resource-scoped; no VM lock).
LEASE_KINDS = frozenset(
    {
        OperationKind.DEVICE_CLAIM,
        OperationKind.DEVICE_RELEASE,
        OperationKind.DEVICE_TRANSFER,
    }
)

# The cross-backend attach/detach transaction kinds (VM-scoped).
ATTACH_KINDS = frozenset(
    {
        OperationKind.VM_ATTACH_DEVICE,
        OperationKind.VM_DETACH_DEVICE,
    }
)


class AttachRegistry:
    """Tracks live (vm, device) → lease bindings recorded by a successful attach.

    ``detach_device`` needs the held ``LeaseId`` to release after the
    hypervisor unplugs the device. The agent normally persists this for audit;
    in-process we keep a thread-safe map. Keyed on ``(VmId, DeviceId)``.
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._by_pair: dict[tuple[VmId, DeviceId], LeaseId] = {}

    def record(self, vm_id: VmId, device_id: DeviceId, lease: LeaseId) -> None:
        with self._lock:
            self._by_pair[(vm_id, device_id)] = lease

    def pop(self, vm_id: VmId, device_id: DeviceId) -> LeaseId | None:
        with self._lock:
            return self._by_pair.pop((vm_id, device_id), None)

    def get(self, vm_id: VmId, device_id: DeviceId) -> LeaseId | None:
        with self._lock:
            return self._by_pair.get((vm_id, device_id))


def dispatch_lease(
    device: DeviceBackend | None,
    kind: OperationKind,
    payload: object,
    token: CancellationToken | None,
) -> object:
    """Route a bare lease ``kind`` → ``DeviceBackend`` call; return the result."""
    if device is None:
        raise TransportError(
            f"operation {kind.value!r} needs a DeviceBackend but none was injected"
        )

    if kind is OperationKind.DEVICE_CLAIM:
        return device.claim(_as(payload, ClaimDevicePayload).device_id)

    if kind is OperationKind.DEVICE_RELEASE:
        device.release(_as(payload, ReleaseDevicePayload).lease_id)
        return None

    if kind is OperationKind.DEVICE_TRANSFER:
        from mgf.vm.core.ids import HostId

        p = _as(payload, TransferLeasePayload)
        return device.transfer_lease(p.lease_id, HostId(p.dest_host), token=token)

    raise TransportError(f"device handler does not own {kind.value!r}")


def dispatch_attach(
    device: DeviceBackend | None,
    hypervisor: HypervisorBackend,
    registry: AttachRegistry,
    kind: OperationKind,
    payload: object,
) -> object:
    """Run the attach (or detach) transaction. Returns ``None``.

    See the module docstring for the rollback / needs-human contract.
    """
    if device is None:
        raise TransportError(
            f"operation {kind.value!r} needs a DeviceBackend but none was injected"
        )

    if kind is OperationKind.VM_ATTACH_DEVICE:
        p = _as(payload, VmAttachDevicePayload)
        _attach(device, hypervisor, registry, p)
        return None

    if kind is OperationKind.VM_DETACH_DEVICE:
        q = _as(payload, DetachDevicePayload)
        _detach(device, hypervisor, registry, q)
        return None

    raise TransportError(f"device handler does not own {kind.value!r}")


def _attach(
    device: DeviceBackend,
    hypervisor: HypervisorBackend,
    registry: AttachRegistry,
    p: VmAttachDevicePayload,
) -> None:
    """THE ATTACH TRANSACTION — claim → assignment → attach → record; roll back."""
    domain = DomainId(p.vm_id.value)
    lease = device.claim(p.device_id, mode=ClaimMode.EXCLUSIVE)
    try:
        assignment = device.assignment_for(lease)
        hypervisor.attach_device(domain, assignment, live=p.live)
    except BaseException:
        # ROLL BACK: release the lease so the broker safe-resets the device.
        # Best-effort — a release failure must not mask the original cause, so
        # it is swallowed (the broker's liveness fence is the backstop).
        try:
            device.release(lease)
        except VmManagerError:
            pass
        raise
    # Success: the lease stays held; record it so detach can release it.
    registry.record(p.vm_id, p.device_id, lease)


def _detach(
    device: DeviceBackend,
    hypervisor: HypervisorBackend,
    registry: AttachRegistry,
    p: DetachDevicePayload,
) -> None:
    """Mirror in reverse: hypervisor detach FIRST, then broker release.

    If the hypervisor detach fails the lease is NOT released (releasing would
    ghost a device still presented to the guest). The failure is re-raised as
    ``DeviceDetachFailed`` — the CP marks the op needs-human (``contracts.md``
    §5). The lease binding is left in place so a human / retry can resolve it.
    """
    domain = DomainId(p.vm_id.value)
    lease = registry.get(p.vm_id, p.device_id)
    assignment = device.assignment_for(lease) if lease is not None else None

    try:
        if assignment is not None:
            hypervisor.detach_device(domain, assignment, live=p.live)
        else:
            # No recorded lease (idempotent re-detach, or attached out of band):
            # nothing the SDK can unplug deterministically. Treat as success —
            # the hypervisor's own attach idempotency makes this safe.
            return
    except DeviceDetachFailed:
        raise
    except VmManagerError as exc:
        raise DeviceDetachFailed(
            f"hypervisor detach of {p.device_id} from {p.vm_id} failed; lease "
            f"{lease} NOT released (needs-human): {exc}"
        ) from exc

    # Hypervisor detach succeeded — now safe to release the lease.
    registry.pop(p.vm_id, p.device_id)
    if lease is not None:
        device.release(lease)


_T = TypeVar("_T")


def _as(payload: object, expected: type[_T]) -> _T:
    """Narrow ``payload`` to ``expected`` or raise a typed transport error."""
    if isinstance(payload, expected):
        return payload
    raise TransportError(f"payload {type(payload).__name__!r} is not a {expected.__name__}")
