"""``FakeDeviceBackend`` — in-memory implementation of
:class:`~mgf.vm.core.contracts.DeviceBackend`.

Simulates the ``mgf-hrb`` broker's behavior: devices have a state
(assignable / leased), claims are exclusive by default, releases
safe-reset, dead-client liveness fencing is simulated via the
:meth:`fence_lease` helper exposed for tests.
"""

from __future__ import annotations

import threading
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING

from mgf.vm.core.capabilities import DeviceCapabilities
from mgf.vm.core.contracts import (
    ClaimMode,
    CompatibilityResult,
    DeviceAssignment,
    DeviceInfo,
    LeaseEvent,
    LeaseInfo,
)
from mgf.vm.core.exceptions import (
    DeviceContended,
    DeviceNotFound,
    LeaseExpired,
    LeaseNotFound,
)
from mgf.vm.core.ids import DeviceId, DomainId, HostId, LeaseId

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator

    from mgf.vm.core.cancellation import CancellationToken


def _utc_now() -> datetime:
    return datetime.now(UTC)


@dataclass
class _SeedDevice:
    """A device the fake "discovers" on the host."""

    device_id: DeviceId
    kind: str
    vendor: str
    product: str
    address: str
    """Bus-specific address (PCI BDF, USB port, …)."""


@dataclass
class _LeaseRecord:
    lease: LeaseId
    device_id: DeviceId
    claimant: str
    mode: ClaimMode
    held_since: datetime
    expired: bool = False
    """Set by :meth:`fence_lease` to simulate liveness-fenced loss."""


class FakeDeviceBackend:
    """In-memory device-passthrough adapter — D-F3 first-class fake.

    Constructed with a list of "discovered" devices. Tests can also
    inject more via :meth:`_register_device`.
    """

    def __init__(
        self,
        *,
        clock: Callable[[], datetime] = _utc_now,
        claimant: str = "fake-claimant",
        seed_devices: list[_SeedDevice] | None = None,
    ) -> None:
        self._clock = clock
        self._claimant = claimant
        self._lock = threading.RLock()
        self._devices: dict[DeviceId, _SeedDevice] = {}
        self._leases: dict[LeaseId, _LeaseRecord] = {}
        # device_id → LeaseId for fast contention checks.
        self._device_lease: dict[DeviceId, LeaseId] = {}
        self._closed = False
        if seed_devices is not None:
            for dev in seed_devices:
                self._devices[dev.device_id] = dev

    @classmethod
    def with_default_devices(
        cls,
        *,
        clock: Callable[[], datetime] = _utc_now,
        claimant: str = "fake-claimant",
    ) -> FakeDeviceBackend:
        """Convenience: a fake with three pre-discovered devices.

        Useful for smoke tests that just want a working broker
        without enumerating seed_devices.
        """
        return cls(
            clock=clock,
            claimant=claimant,
            seed_devices=[
                _SeedDevice(
                    device_id=DeviceId("dev:pci:0000:01:00.0"),
                    kind="pci",
                    vendor="10de",
                    product="1eb1",
                    address="0000:01:00.0",
                ),
                _SeedDevice(
                    device_id=DeviceId("dev:usb:1-2.3"),
                    kind="usb",
                    vendor="046d",
                    product="c52b",
                    address="1-2.3",
                ),
                _SeedDevice(
                    device_id=DeviceId("dev:pci:0000:02:00.0"),
                    kind="pci",
                    vendor="8086",
                    product="1521",
                    address="0000:02:00.0",
                ),
            ],
        )

    @property
    def closed(self) -> bool:
        return self._closed

    @property
    def claimant(self) -> str:
        return self._claimant

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

    def close(self) -> None:
        with self._lock:
            self._closed = True

    def __enter__(self) -> FakeDeviceBackend:
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()

    # --- queries ---------------------------------------------------------

    def list_assignable(self) -> list[DeviceInfo]:
        """Devices NOT currently leased (or leased by other clients)
        that this client could claim."""
        with self._lock:
            assignable = []
            for dev in self._devices.values():
                lease_id = self._device_lease.get(dev.device_id)
                holder = self._leases.get(lease_id) if lease_id else None
                # A device is assignable if not leased OR if we hold it.
                is_free = holder is None or holder.claimant == self._claimant
                assignable.append(
                    DeviceInfo(
                        device_id=dev.device_id,
                        kind=dev.kind,
                        vendor=dev.vendor,
                        product=dev.product,
                        address=dev.address,
                        assignable=is_free,
                    )
                )
            return assignable

    def list_leased(self) -> list[LeaseInfo]:
        with self._lock:
            return [
                LeaseInfo(
                    lease=r.lease,
                    device_id=r.device_id,
                    claimant=r.claimant,
                    mode=r.mode,
                    held_since=r.held_since,
                )
                for r in self._leases.values()
                if not r.expired
            ]

    def get_device(self, d: DeviceId) -> DeviceInfo:
        with self._lock:
            dev = self._require_device(d)
            lease_id = self._device_lease.get(d)
            holder = self._leases.get(lease_id) if lease_id else None
            is_free = holder is None or holder.claimant == self._claimant
            return DeviceInfo(
                device_id=dev.device_id,
                kind=dev.kind,
                vendor=dev.vendor,
                product=dev.product,
                address=dev.address,
                assignable=is_free,
            )

    def check_compatibility(
        self,
        domain: DomainId,
        device: DeviceId,
    ) -> CompatibilityResult:
        with self._lock:
            self._require_device(device)
            # The fake assumes all devices are compatible with all
            # domains. Real adapters do CPU/IOMMU checks here.
            return CompatibilityResult(compatible=True)

    # --- lease lifecycle -------------------------------------------------

    def claim(
        self,
        device: DeviceId,
        *,
        mode: ClaimMode = ClaimMode.EXCLUSIVE,
        ttl: float | None = None,
    ) -> LeaseId:
        with self._lock:
            self._require_device(device)
            existing_lease_id = self._device_lease.get(device)
            if existing_lease_id is not None:
                existing = self._leases.get(existing_lease_id)
                if existing is not None and not existing.expired:
                    if existing.claimant == self._claimant and existing.mode == mode:
                        # Same client, same mode → idempotent.
                        return existing_lease_id
                    if existing.claimant != self._claimant:
                        msg = f"device {device} is leased by {existing.claimant!r}"
                        raise DeviceContended(msg)
                    # Same claimant but different mode — refuse; release first.
                    msg = (
                        f"device {device} already leased by us with a different mode; release first"
                    )
                    raise DeviceContended(msg)
            new_id = LeaseId(f"fake-lease-{uuid.uuid4()}")
            self._leases[new_id] = _LeaseRecord(
                lease=new_id,
                device_id=device,
                claimant=self._claimant,
                mode=mode,
                held_since=self._clock(),
            )
            self._device_lease[device] = new_id
            return new_id

    def assignment_for(self, lease: LeaseId) -> DeviceAssignment:
        with self._lock:
            record = self._leases.get(lease)
            if record is None:
                msg = f"no lease {lease}"
                raise LeaseNotFound(msg)
            if record.expired:
                msg = f"lease {lease} expired (liveness fence)"
                raise LeaseExpired(msg)
            dev = self._require_device(record.device_id)
            return DeviceAssignment(
                lease=lease,
                device_id=record.device_id,
                host_handle=dev.address,
            )

    def release(self, lease: LeaseId) -> None:
        with self._lock:
            record = self._leases.get(lease)
            if record is None:
                # Idempotent.
                return
            self._device_lease.pop(record.device_id, None)
            del self._leases[lease]

    def transfer_lease(
        self,
        lease: LeaseId,
        dest_host: HostId,
        *,
        token: CancellationToken | None = None,
    ) -> LeaseId:
        """Simulated lease handoff.

        In the real broker this coordinates a release-on-source +
        claim-on-destination atomic transaction. The fake just
        rotates the lease ID to simulate the cross-host new identity.
        """
        if token is not None:
            token.raise_if_cancelled()
        with self._lock:
            record = self._leases.get(lease)
            if record is None:
                msg = f"no lease {lease} to transfer"
                raise LeaseNotFound(msg)
            if record.expired:
                msg = f"lease {lease} expired before transfer"
                raise LeaseExpired(msg)
            new_lease_id = LeaseId(f"fake-lease-{uuid.uuid4()}")
            # Rotate identity; keep state under the new ID.
            self._leases[new_lease_id] = _LeaseRecord(
                lease=new_lease_id,
                device_id=record.device_id,
                claimant=record.claimant,
                mode=record.mode,
                held_since=self._clock(),
            )
            self._device_lease[record.device_id] = new_lease_id
            del self._leases[lease]
            return new_lease_id

    # --- events ----------------------------------------------------------

    def watch_leases(self) -> Iterator[LeaseEvent]:
        """Empty event stream for this fake.

        Subclass to implement actual streaming for tests that need it.
        """
        return iter([])

    # --- capabilities ----------------------------------------------------

    def capabilities(self) -> DeviceCapabilities:
        return DeviceCapabilities(
            backend="fake-hrb",
            backend_version="0.3.0",
            supports_pci=True,
            supports_usb=True,
            supports_gpu_mdev=False,
            supports_sr_iov=True,
            supports_lease_transfer=True,
            fenced_by_liveness=True,
            capability_scoped=True,
        )

    # --- test helpers (not part of the Protocol) -------------------------

    def _register_device(self, dev: _SeedDevice) -> None:
        """Add a device to the discovered set (tests inject hardware)."""
        with self._lock:
            self._devices[dev.device_id] = dev

    def fence_lease(self, lease: LeaseId) -> None:
        """Simulate the broker's liveness fence — the lease expires.

        Subsequent ``assignment_for(lease)`` raises ``LeaseExpired``.
        Used by CONF-D-04 (client death releases lease).
        """
        with self._lock:
            record = self._leases.get(lease)
            if record is None:
                return
            record.expired = True
            # The broker would now safe-reset and free the device.
            self._device_lease.pop(record.device_id, None)

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

    def _require_device(self, d: DeviceId) -> _SeedDevice:
        dev = self._devices.get(d)
        if dev is None:
            msg = f"no device {d}"
            raise DeviceNotFound(msg)
        return dev


__all__ = ["FakeDeviceBackend"]
