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

Models pools as named namespaces; volumes live inside pools.
Snapshots reference volumes; volumes with snapshots refuse delete
(reference-counted GC per ``contracts.md`` §4).
"""

from __future__ import annotations

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

from mgf.vm.core.capabilities import StorageCapabilities
from mgf.vm.core.contracts import (
    PoolInfo,
    PoolSpec,
    VolumeSpec,
    _SnapshotInfoContract,
    _VolumeInfoContract,
)
from mgf.vm.core.exceptions import (
    FormatUnsupported,
    InsufficientCapacity,
    PoolExists,
    PoolNotEmpty,
    PoolNotFound,
    ResizeUnsupported,
    SnapshotNotFound,
    VolumeExists,
    VolumeHasSnapshots,
    VolumeInUse,
    VolumeNotFound,
)
from mgf.vm.core.ids import SnapshotId, VolumeId

if TYPE_CHECKING:
    from collections.abc import Callable

    from mgf.vm.core.cancellation import CancellationToken


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


@dataclass
class _PoolRecord:
    name: str
    pool_type: str
    target_path: str
    capacity_bytes: int
    allocation_bytes: int = 0
    active: bool = True
    volume_ids: set[VolumeId] = field(default_factory=set)


@dataclass
class _VolumeRecord:
    volume_id: VolumeId
    pool: str
    name: str
    size_bytes: int
    format: str
    backing_file: str | None
    snapshot_ids: set[SnapshotId] = field(default_factory=set)
    in_use: bool = False
    """Set by tests / by the hypervisor fake when a VM is attached."""


class FakeStorageBackend:
    """In-memory storage adapter — D-F3 first-class fake.

    Supported formats: ``qcow2``, ``raw``. Other formats raise
    ``FormatUnsupported`` — lets adapter conformance exercise the
    capability-gating.

    Pool defaults: a single ``default`` pool of 100 GiB. Tests can
    add more via :meth:`create_pool`.
    """

    SUPPORTED_FORMATS: frozenset[str] = frozenset({"qcow2", "raw"})
    SUPPORTED_POOL_TYPES: frozenset[str] = frozenset({"dir", "lvm"})

    def __init__(
        self,
        *,
        clock: Callable[[], datetime] = _utc_now,
        default_pool_capacity_bytes: int = 100 * (1 << 30),
        supports_shrink: bool = False,
    ) -> None:
        self._clock = clock
        self._supports_shrink = supports_shrink
        self._lock = threading.RLock()
        self._pools: dict[str, _PoolRecord] = {}
        self._volumes: dict[VolumeId, _VolumeRecord] = {}
        self._snapshots: dict[SnapshotId, _SnapshotInfoContract] = {}
        self._closed = False
        # Seed a default pool so smoke tests don't have to create one.
        self._pools["default"] = _PoolRecord(
            name="default",
            pool_type="dir",
            target_path="/var/lib/fake/pool",
            capacity_bytes=default_pool_capacity_bytes,
        )

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

    # --- connection lifecycle --------------------------------------------

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

    def __enter__(self) -> FakeStorageBackend:
        return self

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

    # --- pools -----------------------------------------------------------

    def list_pools(self) -> list[PoolInfo]:
        with self._lock:
            return [self._pool_info(p) for p in self._pools.values()]

    def get_pool(self, name: str) -> PoolInfo:
        with self._lock:
            return self._pool_info(self._require_pool(name))

    def refresh_pool(self, name: str) -> PoolInfo:
        # The fake has no out-of-band volume creation to rescan, so a
        # refresh is observably a ``get_pool`` (D-G20).
        with self._lock:
            return self._pool_info(self._require_pool(name))

    def create_pool(self, spec: PoolSpec) -> None:
        with self._lock:
            existing = self._pools.get(spec.name)
            if existing is not None:
                if (
                    existing.pool_type == spec.pool_type
                    and existing.target_path == spec.target_path
                ):
                    # Same spec → idempotent.
                    return
                msg = f"pool {spec.name!r} exists with a different spec"
                raise PoolExists(msg)
            self._pools[spec.name] = _PoolRecord(
                name=spec.name,
                pool_type=spec.pool_type,
                target_path=spec.target_path,
                capacity_bytes=100 * (1 << 30),
            )

    def delete_pool(self, name: str, *, force: bool = False) -> None:
        with self._lock:
            pool = self._pools.get(name)
            if pool is None:
                # Idempotent.
                return
            if pool.volume_ids and not force:
                msg = f"pool {name!r} contains {len(pool.volume_ids)} volumes"
                raise PoolNotEmpty(msg)
            # force: cascade delete.
            for vol_id in list(pool.volume_ids):
                self._delete_volume_unlocked(vol_id)
            del self._pools[name]

    # --- volumes ---------------------------------------------------------

    def list_volumes(self, pool: str) -> list[_VolumeInfoContract]:
        with self._lock:
            pool_rec = self._require_pool(pool)
            return [
                self._volume_info(self._volumes[v])
                for v in pool_rec.volume_ids
                if v in self._volumes
            ]

    def get_volume(self, v: VolumeId) -> _VolumeInfoContract:
        with self._lock:
            return self._volume_info(self._require_volume(v))

    def create_volume(self, spec: VolumeSpec) -> VolumeId:
        with self._lock:
            if spec.format not in self.SUPPORTED_FORMATS:
                msg = f"format {spec.format!r} not supported (supported: {sorted(self.SUPPORTED_FORMATS)})"
                raise FormatUnsupported(msg)
            pool = self._require_pool(spec.pool)
            # Idempotency: check for existing volume in same pool with same name.
            for vol_id in pool.volume_ids:
                existing = self._volumes.get(vol_id)
                if existing is not None and existing.name == spec.name:
                    if (
                        existing.size_bytes == spec.size_bytes
                        and existing.format == spec.format
                        and existing.backing_file == spec.backing_file
                    ):
                        return vol_id
                    msg = f"volume {spec.name!r} exists in pool {spec.pool!r} with different spec"
                    raise VolumeExists(msg)
            # Capacity check.
            if pool.allocation_bytes + spec.size_bytes > pool.capacity_bytes:
                msg = f"pool {spec.pool!r} insufficient capacity"
                raise InsufficientCapacity(msg)
            new_id = VolumeId(f"fake-vol-{uuid.uuid4()}")
            self._volumes[new_id] = _VolumeRecord(
                volume_id=new_id,
                pool=spec.pool,
                name=spec.name,
                size_bytes=spec.size_bytes,
                format=spec.format,
                backing_file=spec.backing_file,
            )
            pool.volume_ids.add(new_id)
            pool.allocation_bytes += spec.size_bytes
            return new_id

    def resize_volume(self, v: VolumeId, new_size_bytes: int) -> None:
        with self._lock:
            record = self._require_volume(v)
            if new_size_bytes == record.size_bytes:
                # Idempotent.
                return
            if new_size_bytes < record.size_bytes and not self._supports_shrink:
                msg = "this storage backend does not support shrinking"
                raise ResizeUnsupported(msg)
            pool = self._require_pool(record.pool)
            delta = new_size_bytes - record.size_bytes
            if pool.allocation_bytes + delta > pool.capacity_bytes:
                msg = f"pool {record.pool!r} insufficient capacity for resize"
                raise InsufficientCapacity(msg)
            record.size_bytes = new_size_bytes
            pool.allocation_bytes += delta

    def delete_volume(self, v: VolumeId) -> None:
        with self._lock:
            self._delete_volume_unlocked(v)

    def _delete_volume_unlocked(self, v: VolumeId) -> None:
        record = self._volumes.get(v)
        if record is None:
            # Idempotent.
            return
        if record.in_use:
            msg = f"volume {record.name!r} is in use"
            raise VolumeInUse(msg)
        if record.snapshot_ids:
            msg = f"volume {record.name!r} has {len(record.snapshot_ids)} snapshots"
            raise VolumeHasSnapshots(msg)
        pool = self._pools.get(record.pool)
        if pool is not None:
            pool.volume_ids.discard(v)
            pool.allocation_bytes -= record.size_bytes
        del self._volumes[v]

    # --- snapshots -------------------------------------------------------

    def snapshot_volume(self, v: VolumeId, name: str) -> SnapshotId:
        with self._lock:
            record = self._require_volume(v)
            for snap_id in record.snapshot_ids:
                snap = self._snapshots.get(snap_id)
                if snap is not None and snap.name == name:
                    return snap_id
            new_id = SnapshotId(f"fake-vsnap-{uuid.uuid4()}")
            self._snapshots[new_id] = _SnapshotInfoContract(
                snapshot_id=new_id,
                parent=v,
                name=name,
                created_at=self._clock(),
                has_memory=False,
            )
            record.snapshot_ids.add(new_id)
            return new_id

    def list_snapshots(self, v: VolumeId) -> list[_SnapshotInfoContract]:
        with self._lock:
            record = self._require_volume(v)
            return [self._snapshots[s] for s in record.snapshot_ids if s in self._snapshots]

    def delete_snapshot(self, s: SnapshotId) -> None:
        with self._lock:
            snap = self._snapshots.get(s)
            if snap is None:
                return  # idempotent
            # This fake doesn't model the libvirt backing-chain
            # parent/child snapshot graph; real adapters check
            # for derived snapshots here and raise SnapshotHasChildren.
            parent_volume_id = snap.parent  # VolumeId | DomainId
            if isinstance(parent_volume_id, VolumeId):
                parent_volume = self._volumes.get(parent_volume_id)
                if parent_volume is not None:
                    parent_volume.snapshot_ids.discard(s)
            del self._snapshots[s]

    def revert_volume_to_snapshot(
        self,
        v: VolumeId,
        s: SnapshotId,
        *,
        token: CancellationToken | None = None,
    ) -> None:
        if token is not None:
            token.raise_if_cancelled()
        with self._lock:
            self._require_volume(v)
            snap = self._snapshots.get(s)
            if snap is None or snap.parent != v:
                msg = f"snapshot {s} not found on volume {v}"
                raise SnapshotNotFound(msg)
            # In-memory revert is observably a no-op against the
            # volume's recorded fields; the snapshot itself records
            # what would be restored.

    # --- format conversion -----------------------------------------------

    def convert_format(
        self,
        src: VolumeId,
        dst_format: str,
        dst_spec: VolumeSpec,
        *,
        token: CancellationToken | None = None,
    ) -> VolumeId:
        if token is not None:
            token.raise_if_cancelled()
        with self._lock:
            self._require_volume(src)
            if dst_format not in self.SUPPORTED_FORMATS:
                msg = f"target format {dst_format!r} not supported"
                raise FormatUnsupported(msg)
            # Convert = create a new volume with the dst spec.
            return self.create_volume(dst_spec)

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

    def capabilities(self) -> StorageCapabilities:
        return StorageCapabilities(
            engine="fake",
            engine_version="0.3.0",
            supports_pool_types=self.SUPPORTED_POOL_TYPES,
            supports_formats=self.SUPPORTED_FORMATS,
            supports_shrink=self._supports_shrink,
            supports_thin_provisioning=True,
        )

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

    def _require_pool(self, name: str) -> _PoolRecord:
        pool = self._pools.get(name)
        if pool is None:
            msg = f"no pool {name!r}"
            raise PoolNotFound(msg)
        return pool

    def _require_volume(self, v: VolumeId) -> _VolumeRecord:
        record = self._volumes.get(v)
        if record is None:
            msg = f"no volume {v}"
            raise VolumeNotFound(msg)
        return record

    @staticmethod
    def _pool_info(p: _PoolRecord) -> PoolInfo:
        return PoolInfo(
            name=p.name,
            pool_type=p.pool_type,
            target_path=p.target_path,
            capacity_bytes=p.capacity_bytes,
            allocation_bytes=p.allocation_bytes,
            available_bytes=max(p.capacity_bytes - p.allocation_bytes, 0),
            active=p.active,
        )

    def _volume_info(self, record: _VolumeRecord) -> _VolumeInfoContract:
        pool = self._pools.get(record.pool)
        target = pool.target_path if pool is not None else ""
        return _VolumeInfoContract(
            volume_id=record.volume_id,
            pool=record.pool,
            name=record.name,
            path=f"{target}/{record.name}",
            size_bytes=record.size_bytes,
            allocation_bytes=record.size_bytes,
            format=record.format,
        )
