# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Process-local ownership for immutable mmap-backed host tensors."""

from __future__ import annotations

import contextlib
import threading
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import NoReturn, SupportsIndex

import torch

from .identity import WeightArtifactIdentity
from .manifest import WeightManifest


@dataclass(frozen=True)
class MappedHostRegion:
    """One process-local virtual range belonging to an artifact payload file."""

    file_name: str
    address: int
    length: int
    read_only: bool = True

    def __post_init__(self) -> None:
        if not self.file_name or self.address < 0 or self.length < 0:
            raise ValueError("mapped host region contains invalid fields")


@dataclass(frozen=True)
class LeaseProvenance:
    resolution_id: str
    identity_digest: str
    artifact_content_sha256: str
    domain_uuid: str
    publication_id: str
    resolution_source: str
    artifact_origin: str
    producer_id: str | None
    provider_id: str | None
    validation_level: str
    backing_fingerprint: str


class HostWeightLease:
    """Stable tensor views and resources retained until transport teardown."""

    def __init__(
        self,
        *,
        identity: WeightArtifactIdentity,
        manifest: WeightManifest,
        tensors: Mapping[str, torch.Tensor],
        mapped_regions: tuple[MappedHostRegion, ...],
        provenance: LeaseProvenance,
        resource_closers: tuple[Callable[[], None], ...],
        release_artifact_lock: Callable[[], None],
    ) -> None:
        if manifest.identity.canonical_bytes != identity.canonical_bytes:
            raise ValueError("lease manifest identity differs from the requested identity")
        if set(tensors) != {entry.name for entry in manifest.tensors}:
            raise ValueError("lease tensor names differ from the validated manifest")
        entries = {entry.name: entry for entry in manifest.tensors}
        for name, tensor in tensors.items():
            entry = entries[name]
            if (
                not isinstance(tensor, torch.Tensor)
                or tensor.device.type != "cpu"
                or tuple(tensor.shape) != entry.shape
                or str(tensor.dtype) != entry.dtype
                or tuple(tensor.stride()) != entry.stride
                or tensor.numel() * tensor.element_size() != entry.nbytes
            ):
                raise ValueError(f"lease tensor {name!r} differs from the validated host metadata")
        if provenance.identity_digest != identity.key:
            raise ValueError("lease provenance identity digest differs from the requested identity")
        if provenance.artifact_content_sha256 != manifest.artifact_content_sha256:
            raise ValueError("lease provenance content digest differs from the validated manifest")
        if not isinstance(mapped_regions, tuple) or any(
            not isinstance(region, MappedHostRegion) for region in mapped_regions
        ):
            raise ValueError("lease mapped regions must be an immutable tuple")
        if not isinstance(resource_closers, tuple) or any(not callable(closer) for closer in resource_closers):
            raise ValueError("lease resource closers must be an immutable tuple of callables")
        if not callable(release_artifact_lock):
            raise ValueError("lease artifact lock release must be callable")
        self.identity = identity
        self.manifest = manifest
        self._tensor_storage = dict(tensors)
        self.tensors = MappingProxyType(self._tensor_storage)
        self.mapped_regions = mapped_regions
        self.provenance = provenance
        self._resource_closers = resource_closers
        self._release_artifact_lock = release_artifact_lock
        self._state_lock = threading.RLock()
        self._closed = False

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

    def close(self) -> None:
        with self._state_lock:
            if self._closed:
                return
            self._closed = True

            errors: list[BaseException] = []
            self._tensor_storage.clear()
            for closer in reversed(self._resource_closers):
                try:
                    closer()
                except BaseException as exc:  # pragma: no cover - best-effort teardown aggregation
                    errors.append(exc)
            try:
                self._release_artifact_lock()
            except BaseException as exc:  # pragma: no cover - best-effort teardown aggregation
                errors.append(exc)
            if errors:
                raise RuntimeError(f"failed to close {len(errors)} host weight lease resources") from errors[0]

    def __enter__(self) -> HostWeightLease:
        if self.closed:
            raise RuntimeError("cannot re-enter a closed HostWeightLease")
        return self

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

    def __reduce_ex__(self, _protocol: SupportsIndex) -> NoReturn:
        raise TypeError("HostWeightLease is process-local and cannot be serialized")

    def __del__(self) -> None:
        with contextlib.suppress(Exception):
            self.close()


__all__ = ["HostWeightLease", "LeaseProvenance", "MappedHostRegion"]