# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Store-scoped streaming safetensors publication."""

from __future__ import annotations

import contextlib
import errno
import hashlib
import json
import os
import struct
import uuid
from collections.abc import Callable
from pathlib import Path, PurePosixPath
from types import TracebackType

import torch
from safetensors import safe_open

from ..identity import WeightArtifactIdentity, canonical_json
from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest
from ..protocols import ProductionMetadata, TensorWriteSpec

_HASH_CHUNK_BYTES = 8 * 1024**2
_MANIFEST_FILE = "manifest.json"
_READY_FILE = "READY.json"

_SAFETENSORS_DTYPES: dict[torch.dtype, str] = {
    torch.bool: "BOOL",
    torch.uint8: "U8",
    torch.int8: "I8",
    torch.int16: "I16",
    torch.int32: "I32",
    torch.int64: "I64",
    torch.float16: "F16",
    torch.bfloat16: "BF16",
    torch.float32: "F32",
    torch.float64: "F64",
}
for _torch_name, _safe_name in (
    ("float8_e4m3fn", "F8_E4M3"),
    ("float8_e5m2", "F8_E5M2"),
    ("float8_e4m3fnuz", "F8_E4M3FNUZ"),
    ("float8_e5m2fnuz", "F8_E5M2FNUZ"),
):
    if (_dtype := getattr(torch, _torch_name, None)) is not None:
        _SAFETENSORS_DTYPES[_dtype] = _safe_name


class ArtifactWriterError(RuntimeError):
    pass


def _safe_relative_name(value: str) -> str:
    path = PurePosixPath(value)
    if not value or path.is_absolute() or len(path.parts) != 1 or path.name != value or value in {".", ".."}:
        raise ArtifactWriterError(f"artifact payload name must be one safe relative basename: {value!r}")
    if not value.endswith(".safetensors"):
        raise ArtifactWriterError(f"artifact payload must use the .safetensors format: {value!r}")
    return value


def _pwrite_all(fd: int, data: bytes | bytearray | memoryview, offset: int) -> None:
    view = memoryview(data).cast("B")
    written = 0
    while written < len(view):
        count = os.pwrite(fd, view[written:], offset + written)
        if count <= 0:
            raise OSError(errno.EIO, "short pwrite while publishing host weights")
        written += count


def _fsync_directory(path: Path) -> None:
    fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)


def _file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(_HASH_CHUNK_BYTES):
            digest.update(chunk)
    return digest.hexdigest()


def _tensor_sha256(tensor: torch.Tensor) -> str:
    digest = hashlib.sha256()
    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
        digest.update(memoryview(chunk.numpy()))
    return digest.hexdigest()


class _Coverage:
    def __init__(self, size: int) -> None:
        self._size = size
        self._ranges: list[tuple[int, int]] = []

    def add(self, start: int, end: int) -> None:
        if start < 0 or end < start or end > self._size:
            raise ArtifactWriterError(f"write range [{start}, {end}) exceeds tensor storage {self._size}")
        if start == end:
            return
        for existing_start, existing_end in self._ranges:
            if start < existing_end and existing_start < end:
                raise ArtifactWriterError(f"overlapping tensor write ranges [{start}, {end}) and {self._ranges}")
        self._ranges.append((start, end))

    def complete(self) -> bool:
        if self._size == 0:
            return True
        cursor = 0
        for start, end in sorted(self._ranges):
            if start != cursor:
                return False
            cursor = end
        return cursor == self._size


class StreamingSafeTensorFile:
    """One preallocated safetensors file filled through bounded CPU writes."""

    def __init__(
        self,
        owner: FilesystemArtifactWriter,
        relative_name: str,
        specs: tuple[TensorWriteSpec, ...],
    ) -> None:
        self._owner = owner
        self.relative_name = _safe_relative_name(relative_name)
        if not specs:
            raise ArtifactWriterError("a tensor payload file must contain at least one tensor")
        by_name = {spec.name: spec for spec in specs}
        if len(by_name) != len(specs):
            raise ArtifactWriterError("tensor payload contains duplicate storage keys")
        unsupported = [spec.dtype for spec in specs if spec.dtype not in _SAFETENSORS_DTYPES]
        if unsupported:
            raise ArtifactWriterError(f"unsupported safetensors dtypes: {unsupported}")
        self.specs = tuple(sorted(specs, key=lambda spec: spec.name))
        self._specs = {spec.name: spec for spec in self.specs}
        self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs}
        self._closed = False

        offsets: dict[str, int] = {}
        cursor = 0
        for spec in self.specs:
            offsets[spec.name] = cursor
            cursor += spec.nbytes
        header: dict[str, object] = {"__metadata__": {"format": "pt", "host_weight_runtime": "1"}}
        for spec in self.specs:
            start = offsets[spec.name]
            header[spec.name] = {
                "dtype": _SAFETENSORS_DTYPES[spec.dtype],
                "shape": list(spec.shape),
                "data_offsets": [start, start + spec.nbytes],
            }
        encoded_header = json.dumps(header, sort_keys=True, separators=(",", ":")).encode("utf-8")
        encoded_header += b" " * (-len(encoded_header) % 8)
        self._payload_offset = 8 + len(encoded_header)
        self._offsets = offsets
        self._path = owner.staging_dir / self.relative_name
        total_size = self._payload_offset + cursor
        owner._reserve_file(total_size)
        self._fd = os.open(self._path, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_CLOEXEC, 0o600)
        try:
            if hasattr(os, "posix_fallocate"):
                try:
                    os.posix_fallocate(self._fd, 0, total_size)
                except OSError as exc:
                    if exc.errno not in (errno.EINVAL, errno.ENOSYS, errno.EOPNOTSUPP):
                        raise
                    os.ftruncate(self._fd, total_size)
            else:  # pragma: no cover - Linux production path has posix_fallocate
                os.ftruncate(self._fd, total_size)
            _pwrite_all(self._fd, struct.pack("<Q", len(encoded_header)), 0)
            _pwrite_all(self._fd, encoded_header, 8)
        except BaseException:
            os.close(self._fd)
            self._path.unlink(missing_ok=True)
            owner._forget_reserved_file(total_size)
            raise
        owner._opened(self)

    @staticmethod
    def _bytes(tensor: torch.Tensor) -> memoryview:
        if tensor.device.type != "cpu" or tensor.is_meta:
            raise ArtifactWriterError("artifact producers must provide materialized CPU tensor chunks")
        contiguous = tensor.detach().contiguous()
        return memoryview(contiguous.view(torch.uint8).numpy()).cast("B")

    def write_tensor(self, name: str, tensor: torch.Tensor) -> None:
        spec = self._get_spec(name)
        if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype:
            raise ArtifactWriterError(
                f"tensor {name!r} differs from its declared shape/dtype: "
                f"got {tuple(tensor.shape)}/{tensor.dtype}, expected {spec.shape}/{spec.dtype}"
            )
        self._write(name, 0, self._bytes(tensor))

    def write_rows(self, name: str, start_row: int, tensor: torch.Tensor) -> None:
        spec = self._get_spec(name)
        if not spec.shape:
            if start_row != 0 or tuple(tensor.shape) != ():
                raise ArtifactWriterError(f"scalar tensor {name!r} only accepts one complete scalar write")
            self.write_tensor(name, tensor)
            return
        if tensor.dtype != spec.dtype or tuple(tensor.shape[1:]) != spec.shape[1:]:
            raise ArtifactWriterError(f"row chunk for {name!r} differs from its declared dtype or trailing shape")
        row_count = tensor.shape[0]
        if start_row < 0 or start_row + row_count > spec.shape[0]:
            raise ArtifactWriterError(f"row chunk for {name!r} exceeds its declared first dimension")
        row_bytes = spec.nbytes // spec.shape[0] if spec.shape[0] else 0
        self._write(name, start_row * row_bytes, self._bytes(tensor))

    def _get_spec(self, name: str) -> TensorWriteSpec:
        if self._closed:
            raise ArtifactWriterError("cannot write through a closed tensor file")
        try:
            return self._specs[name]
        except KeyError as exc:
            raise ArtifactWriterError(f"tensor {name!r} was not declared in this payload file") from exc

    def _write(self, name: str, byte_offset: int, data: memoryview) -> None:
        coverage = self._coverage[name]
        coverage.add(byte_offset, byte_offset + len(data))
        _pwrite_all(self._fd, data, self._payload_offset + self._offsets[name] + byte_offset)

    def close(self) -> None:
        if self._closed:
            return
        missing = [name for name, coverage in self._coverage.items() if not coverage.complete()]
        if missing:
            self.abort()
            raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}")
        self._closed = True
        try:
            os.fsync(self._fd)
        finally:
            os.close(self._fd)
        self._owner._closed(self)

    def abort(self) -> None:
        if self._closed:
            return
        self._closed = True
        os.close(self._fd)
        size = self._path.stat().st_size if self._path.exists() else 0
        self._path.unlink(missing_ok=True)
        self._owner._aborted(self, size)

    def __enter__(self) -> StreamingSafeTensorFile:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        _traceback: TracebackType | None,
    ) -> None:
        if exc_type is None:
            self.close()
        else:
            self.abort()


class FilesystemArtifactWriter:
    """Scoped writer that cannot publish outside one invisible staging dir."""

    def __init__(
        self,
        staging_dir: Path,
        *,
        capacity_check: Callable[[int, int], None],
    ) -> None:
        self.staging_dir = staging_dir
        self._capacity_check = capacity_check
        self._active: set[StreamingSafeTensorFile] = set()
        self._files: dict[str, tuple[TensorWriteSpec, ...]] = {}
        self._reserved_bytes = 0
        self._finalized = False

    def open_tensor_file(
        self,
        relative_name: str,
        specs: tuple[TensorWriteSpec, ...],
    ) -> StreamingSafeTensorFile:
        if self._finalized:
            raise ArtifactWriterError("cannot add payloads after artifact finalization")
        relative_name = _safe_relative_name(relative_name)
        if relative_name in self._files or any(item.relative_name == relative_name for item in self._active):
            raise ArtifactWriterError(f"artifact payload {relative_name!r} was declared more than once")
        return StreamingSafeTensorFile(self, relative_name, specs)

    def _reserve_file(self, size: int) -> None:
        self._capacity_check(self._reserved_bytes + size, size)
        self._reserved_bytes += size

    def _forget_reserved_file(self, size: int) -> None:
        self._reserved_bytes = max(0, self._reserved_bytes - size)

    def _opened(self, writer: StreamingSafeTensorFile) -> None:
        self._active.add(writer)

    def _closed(self, writer: StreamingSafeTensorFile) -> None:
        self._active.remove(writer)
        self._files[writer.relative_name] = writer.specs

    def _aborted(self, writer: StreamingSafeTensorFile, size: int) -> None:
        self._active.discard(writer)
        self._forget_reserved_file(size)

    def abort(self) -> None:
        for writer in tuple(self._active):
            with contextlib.suppress(Exception):
                writer.abort()

    def finalize(
        self,
        identity: WeightArtifactIdentity,
        metadata: ProductionMetadata,
    ) -> tuple[WeightManifest, PublicationMarker]:
        if self._finalized:
            raise ArtifactWriterError("artifact writer was finalized more than once")
        if self._active:
            raise ArtifactWriterError("artifact producer left tensor payload files open")
        if not self._files:
            raise ArtifactWriterError("artifact producer emitted no tensor payload files")
        if metadata.producer_schema != identity.producer.manifest_schema:
            raise ArtifactWriterError("production metadata does not match the identity's manifest schema")
        if metadata.restorer_schema != identity.producer.restorer_schema:
            raise ArtifactWriterError("production metadata does not match the identity's restorer schema")

        tensor_entries: list[TensorManifestEntry] = []
        file_entries: list[FileManifestEntry] = []
        seen_tensors: set[str] = set()
        for file_name in sorted(self._files):
            path = self.staging_dir / file_name
            os.chmod(path, 0o444)
            with path.open("rb") as handle:
                os.fsync(handle.fileno())
            file_entries.append(
                FileManifestEntry(
                    file_name=file_name,
                    size=path.stat().st_size,
                    sha256=_file_sha256(path),
                )
            )
            expected = {spec.name: spec for spec in self._files[file_name]}
            with safe_open(str(path), framework="pt", device="cpu") as handle:
                if set(handle.keys()) != set(expected):
                    raise ArtifactWriterError(f"published tensor keys differ in {file_name!r}")
                for name in sorted(expected):
                    if name in seen_tensors:
                        raise ArtifactWriterError(f"artifact tensor {name!r} occurs in multiple payload files")
                    seen_tensors.add(name)
                    spec = expected[name]
                    tensor = handle.get_tensor(name)
                    if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype:
                        raise ArtifactWriterError(f"published tensor {name!r} differs from its declared metadata")
                    entry = TensorManifestEntry(
                        name=name,
                        kind=spec.kind,
                        role=spec.role,
                        file_name=file_name,
                        storage_key=name,
                        shape=tuple(tensor.shape),
                        dtype=str(tensor.dtype),
                        stride=tuple(tensor.stride()),
                        nbytes=tensor.numel() * tensor.element_size(),
                        sha256=_tensor_sha256(tensor),
                    )
                    tensor_entries.append(entry)

        ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name))
        content_digest = hashlib.sha256()
        for entry in ordered_tensor_entries:
            content_digest.update(canonical_json(entry.semantic_dict()))

        manifest = WeightManifest(
            schema_version=1,
            artifact_key=identity.key,
            identity=identity,
            producer_schema=metadata.producer_schema,
            restorer_schema=metadata.restorer_schema,
            tensors=ordered_tensor_entries,
            files=tuple(file_entries),
            artifact_content_sha256=content_digest.hexdigest(),
            format_metadata=metadata.format_metadata,
        )
        manifest_bytes = manifest.canonical_bytes
        self._write_metadata_file(_MANIFEST_FILE, manifest_bytes)
        marker = PublicationMarker(
            schema_version=1,
            publication_id=uuid.uuid4().hex,
            manifest_sha256=hashlib.sha256(manifest_bytes).hexdigest(),
            artifact_content_sha256=manifest.artifact_content_sha256,
            integrity_mode="full_checksum",
        )
        self._write_metadata_file(_READY_FILE, marker.canonical_bytes)
        _fsync_directory(self.staging_dir)
        self._finalized = True
        return manifest, marker

    def _write_metadata_file(self, name: str, data: bytes) -> None:
        path = self.staging_dir / name
        fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o600)
        try:
            _pwrite_all(fd, data, 0)
            os.fsync(fd)
            os.fchmod(fd, 0o444)
            os.fsync(fd)
        finally:
            os.close(fd)


__all__ = [
    "ArtifactWriterError",
    "FilesystemArtifactWriter",
    "StreamingSafeTensorFile",
]