# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""IPC utilities for transferring large tensors via POSIX shared memory.

Used by Hop1 (GPU worker <-> scheduler) to avoid pickling large video tensors
through the MessageQueue. Tensors above ``_SHM_TENSOR_THRESHOLD`` are copied
into a named shared-memory segment; only a lightweight metadata dict is
serialised through the queue.
"""

from __future__ import annotations

from collections.abc import Iterator
from contextlib import suppress
from typing import Any, TypeGuard

import numpy as np
import torch
from vllm.logger import init_logger

from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.media import (
    DiffusionMediaOutput,
    FloatVideoConsumer,
    VideoColorModel,
    VideoMediaOutput,
    VideoTensorEncoding,
    VideoTensorLayout,
    VideoTensorSpec,
    VideoTransportConstraints,
    VideoValueRange,
)

_SHM_TENSOR_THRESHOLD = 1_000_000  # 1 MB
DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"
_DIFFUSION_MEDIA_WIRE_TYPE = "diffusion_media_v1"
logger = init_logger(__name__)

# An uncertain CUDA cleanup must not unmap memory still used by the device.
# Retain failed transfers until worker exit and reject further registrations.
_failed_shm_copies: list[tuple[Any, torch.Tensor, torch.Tensor, torch.Stream]] = []

# Sentinel so compute-then-assign packing can tell "field not packed" apart from
# "field packed to None".
_UNSET = object()


class _SharedMemoryLease:
    """Keep an unlinked shared-memory mapping alive while array views use it."""

    def __init__(self, shm: Any, array: np.ndarray) -> None:
        self.shm = shm
        self.__array_interface__ = array.__array_interface__

    def __del__(self) -> None:
        with suppress(BufferError):
            self.shm.close()


def _new_shm_array(shape: tuple[int, ...], dtype: np.dtype) -> tuple[Any, np.ndarray, dict[str, Any]]:
    from multiprocessing import shared_memory

    nbytes = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize
    shm = shared_memory.SharedMemory(create=True, size=nbytes)
    array = np.ndarray(shape, dtype=dtype, buffer=shm.buf[:nbytes])
    handle = {
        "name": shm.name,
        "shape": list(shape),
        "numpy_dtype": str(dtype),
        "nbytes": nbytes,
    }
    return shm, array, handle


def _array_to_shm(array: np.ndarray) -> dict[str, Any]:
    """Copy a contiguous NumPy-compatible array into shared memory."""
    if array.dtype.hasobject:
        raise TypeError("NumPy object arrays cannot be transferred through raw shared memory")

    array = np.ascontiguousarray(array)
    shm, shm_array, handle = _new_shm_array(array.shape, array.dtype)
    np.copyto(shm_array, array)
    del shm_array
    shm.close()
    return handle


def _array_from_shm(handle: dict[str, Any]) -> np.ndarray:
    """Materialize or borrow an array from shared memory and unlink its name."""
    from multiprocessing import shared_memory

    shm = shared_memory.SharedMemory(name=handle["name"])
    array = np.ndarray(
        handle["shape"],
        dtype=np.dtype(handle["numpy_dtype"]),
        buffer=shm.buf[: handle["nbytes"]],
    )
    if handle.get("borrow_on_unpack"):
        # A base owner survives np.asarray(), slicing, and torch.from_numpy().
        # An ndarray subclass attribute does not survive np.asarray().
        borrowed = np.asarray(_SharedMemoryLease(shm, array))
        shm.unlink()
        return borrowed
    try:
        return array.copy()
    finally:
        del array
        shm.close()
        shm.unlink()


def _cuda_registered_shm_copy(tensor: torch.Tensor, d2h_stream: torch.Stream) -> dict[str, Any]:
    """Copy a CUDA tensor directly into a CUDA-registered POSIX SHM segment."""
    if _failed_shm_copies:
        raise RuntimeError("A previous registered SHM transfer failed CUDA cleanup; restart the worker")
    original_dtype = tensor.dtype
    source = tensor.detach()
    # NumPy has no BF16 dtype. Transfer the original bits, without allocating
    # an FP32 accelerator tensor or converting the borrowed CPU result.
    if original_dtype == torch.bfloat16:
        numpy_dtype = np.dtype(np.uint16)
    else:
        numpy_dtype = torch.empty((), dtype=source.dtype).numpy().dtype
    cudart = torch.cuda.cudart()
    shm, shm_array, handle = _new_shm_array(tuple(source.shape), numpy_dtype)
    pointer = shm_array.ctypes.data
    host_tensor = None
    registered = False
    succeeded = False
    try:
        host_tensor = torch.from_numpy(shm_array).view(original_dtype)
        error = cudart.cudaHostRegister(pointer, handle["nbytes"], 0)
        if int(error) != 0:
            raise RuntimeError(f"cudaHostRegister failed: {cudart.cudaGetErrorString(error)}")
        registered = True
        old_stream = torch.accelerator.current_stream()
        torch.accelerator.set_stream(d2h_stream)
        try:
            host_tensor.copy_(source, non_blocking=True)
        finally:
            torch.accelerator.set_stream(old_stream)
        d2h_stream.synchronize()
        error = cudart.cudaHostUnregister(pointer)
        if int(error) != 0:
            raise RuntimeError(f"cudaHostUnregister failed: {cudart.cudaGetErrorString(error)}")
        registered = False
        handle["borrow_on_unpack"] = True
        handle.update(
            {
                "__tensor_shm__": True,
                "torch_dtype": str(original_dtype),
                "torch_dtype_view": original_dtype == torch.bfloat16,
            }
        )
        succeeded = True
        return handle
    finally:
        if registered:
            try:
                # A copy may have been enqueued before a later operation failed.
                d2h_stream.synchronize()
                error = cudart.cudaHostUnregister(pointer)
                if int(error) != 0:
                    raise RuntimeError(f"cudaHostUnregister failed: {cudart.cudaGetErrorString(error)}")
                registered = False
            except Exception:
                assert host_tensor is not None
                _failed_shm_copies.append((shm, host_tensor, source, d2h_stream))
                logger.exception("Retaining registered SHM after CUDA cleanup failure; worker restart required")
        del host_tensor, shm_array
        if not registered:
            shm.close()
        if not succeeded:
            shm.unlink()


def _unlink_shm_handle(handle: dict[str, Any]) -> None:
    """Best-effort removal of a shared-memory segment created during packing."""
    from multiprocessing import shared_memory

    name = handle.get("name")
    if not name:
        return
    try:
        shm = shared_memory.SharedMemory(name=name)
    except FileNotFoundError:
        return
    try:
        shm.close()
    finally:
        try:
            shm.unlink()
        except FileNotFoundError:
            pass


def _unlink_shm_handles(handles: list[dict[str, Any]]) -> None:
    for handle in handles:
        try:
            _unlink_shm_handle(handle)
        except Exception:
            # Cleanup is best-effort; never mask the original packing failure.
            pass


def _tensor_to_shm(
    tensor: torch.Tensor,
    d2h_stream: torch.Stream | None = None,
    created: list[dict[str, Any]] | None = None,
    *,
    enable_registered_shm: bool = False,
) -> dict[str, Any]:
    """Copy a tensor into POSIX shared memory and return a metadata handle.

    The shared memory segment remains alive after this call (the local fd is
    closed, but the segment persists until ``_tensor_from_shm`` unlinks it).

    If *d2h_stream* is provided, the D2H copy uses ``copy_()`` with
    ``pin_memory=True`` on that stream instead of the synchronous ``.cpu()``
    path.  The caller must synchronize *d2h_stream* after all tensors are
    packed.
    """
    original_dtype = tensor.dtype
    if (
        d2h_stream is not None
        and enable_registered_shm
        and tensor.device.type == "cuda"
        # ROCm also uses the "cuda" device type, but not this CUDA runtime API.
        and torch.version.cuda is not None
    ):
        handle = _cuda_registered_shm_copy(tensor, d2h_stream)
        if created is not None:
            created.append(handle)
        return handle
    if d2h_stream is not None:
        # Non-blocking D2H: copy on side stream to pinned CPU memory.
        old_stream = torch.accelerator.current_stream()
        torch.accelerator.set_stream(d2h_stream)
        try:
            t = tensor.detach()
            cpu = torch.empty(t.shape, dtype=t.dtype, pin_memory=True)
            cpu.copy_(t, non_blocking=True)
        finally:
            torch.accelerator.set_stream(old_stream)
        d2h_stream.synchronize()
        tensor = cpu
        # numpy has no bfloat16, so widen on the host after D2H. Never widen on
        # the accelerator before the copy: a payload that already fell back after
        # an fp32 OOM would OOM again on a second device-side fp32 allocation.
        if original_dtype == torch.bfloat16:
            tensor = tensor.to(torch.float32)
    else:
        tensor = tensor.detach().cpu().contiguous()
        if original_dtype == torch.bfloat16:
            tensor = tensor.to(torch.float32)
    handle = _array_to_shm(tensor.numpy())
    handle.update(
        {
            "__tensor_shm__": True,
            "torch_dtype": str(original_dtype),
        }
    )
    if created is not None:
        created.append(handle)
    return handle


def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
    """Reconstruct a tensor from a shared-memory handle and free the segment."""
    tensor = torch.from_numpy(_array_from_shm(handle))
    # Restore the original dtype if it differs from the numpy-compatible
    # dtype used for the SHM transfer (e.g. bfloat16 → float32 → bfloat16).
    torch_dtype_str = handle.get("torch_dtype", "")
    if torch_dtype_str:
        original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
        if original_dtype is not None and tensor.dtype != original_dtype:
            tensor = tensor.view(original_dtype) if handle.get("torch_dtype_view") else tensor.to(original_dtype)
    return tensor


def _pack_tensor_if_large(
    val: torch.Tensor,
    d2h_stream: torch.Stream | None = None,
    created: list[dict[str, Any]] | None = None,
    *,
    enable_registered_shm: bool = False,
) -> torch.Tensor | dict:
    """Replace a tensor with an SHM handle if it exceeds the threshold.

    Batch outputs are split into per-request views that share one storage.
    Pickle serializes a view's whole storage, so a batch of small views costs
    the full batch tensor once per request. Size the decision on the storage
    to keep those views off the wire; packing copies just the view.
    """
    view_bytes = val.nelement() * val.element_size()
    try:
        storage_bytes = val.untyped_storage().nbytes()
    except Exception:
        storage_bytes = view_bytes
    if max(view_bytes, storage_bytes) > _SHM_TENSOR_THRESHOLD:
        return _tensor_to_shm(val, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm)
    return val


def _ndarray_to_shm(array: np.ndarray, created: list[dict[str, Any]] | None = None) -> dict[str, Any]:
    """Copy a contiguous NumPy array into POSIX shared memory."""
    handle = _array_to_shm(array)
    handle["__ndarray_shm__"] = True
    if created is not None:
        created.append(handle)
    return handle


def _ndarray_from_shm(handle: dict[str, Any]) -> np.ndarray:
    """Reconstruct a NumPy array from shared memory and free the segment."""
    return _array_from_shm(handle)


def _pack_value_if_large(
    val: object,
    d2h_stream: torch.Stream | None = None,
    created: list[dict[str, Any]] | None = None,
) -> object:
    """Recursively replace large tensors with SHM handles.

    Walks the container shapes pipelines return as ``DiffusionOutput.output``:
    bare tensors, dicts (e.g. Cosmos3 ``{"image"/"video": ...}``), and
    tuples/lists (e.g. LTX2 ``(video, audio)``). Other values pass
    through unchanged. ``_unpack_if_shm_handle`` must mirror these shapes — keep
    the two in sync.
    """
    if isinstance(val, torch.Tensor):
        return _pack_tensor_if_large(val, d2h_stream=d2h_stream, created=created)
    if isinstance(val, np.ndarray):
        return (
            _ndarray_to_shm(val, created=created)
            if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD
            else val
        )
    if isinstance(val, dict):
        return {key: _pack_value_if_large(value, d2h_stream=d2h_stream, created=created) for key, value in val.items()}
    if isinstance(val, list):
        return [_pack_value_if_large(item, d2h_stream=d2h_stream, created=created) for item in val]
    if isinstance(val, tuple):
        return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream, created=created) for item in val)
    return val


def _unpack_if_shm_handle(val: object) -> object:
    """Reconstruct tensors from SHM handles, mirroring ``_pack_value_if_large``."""
    if isinstance(val, dict) and val.get("__tensor_shm__"):
        return _tensor_from_shm(val)
    if isinstance(val, dict) and val.get("__ndarray_shm__"):
        return _ndarray_from_shm(val)
    if isinstance(val, dict):
        return {key: _unpack_if_shm_handle(value) for key, value in val.items()}
    if isinstance(val, list):
        return [_unpack_if_shm_handle(item) for item in val]
    if isinstance(val, tuple):
        return tuple(_unpack_if_shm_handle(item) for item in val)
    return val


def _pack_diffusion_media(
    media: DiffusionMediaOutput,
    *,
    d2h_stream: torch.Stream | None,
    created: list[dict[str, Any]] | None = None,
    enable_registered_shm: bool = False,
) -> dict[str, Any]:
    media.validate()
    if not media.prepared_for_transport:
        raise ValueError("Diffusion media must be prepared before IPC packing")
    video = media.video
    packed_tensor = _pack_tensor_if_large(
        video.tensor, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm
    )
    if isinstance(packed_tensor, torch.Tensor) and packed_tensor.device.type != "cpu":
        packed_tensor = packed_tensor.detach().cpu().contiguous()
    return {
        "__type__": _DIFFUSION_MEDIA_WIRE_TYPE,
        "schema_version": video.schema_version,
        "prepared_for_transport": True,
        "video": {
            "tensor": packed_tensor,
            "layout": video.spec.layout.value,
            "encoding": video.spec.encoding.value,
            "value_range": video.spec.value_range.value,
            "color_model": video.spec.color_model.value,
            "pending_float_consumers": sorted(consumer.value for consumer in video.constraints.pending_float_consumers),
        },
    }


def _unpack_diffusion_media(packed: dict[str, Any]) -> DiffusionMediaOutput:
    if packed.get("__type__") != _DIFFUSION_MEDIA_WIRE_TYPE:
        raise ValueError("Invalid diffusion media wire payload")
    video_payload = packed.get("video")
    if not isinstance(video_payload, dict):
        raise ValueError("Diffusion media wire payload is missing video metadata")
    tensor = _unpack_if_shm_handle(video_payload.get("tensor"))
    if not isinstance(tensor, torch.Tensor):
        raise ValueError("Diffusion media wire payload contains an invalid video tensor")
    schema_version = packed.get("schema_version")
    if schema_version != 1:
        raise ValueError(f"Unsupported video media schema version: {schema_version}")
    if packed.get("prepared_for_transport") is not True:
        raise ValueError("Diffusion media wire payload was not prepared for transport")
    media = DiffusionMediaOutput(
        video=VideoMediaOutput(
            tensor=tensor,
            spec=VideoTensorSpec(
                layout=VideoTensorLayout(video_payload["layout"]),
                encoding=VideoTensorEncoding(video_payload["encoding"]),
                value_range=VideoValueRange(video_payload["value_range"]),
                color_model=VideoColorModel(video_payload["color_model"]),
            ),
            constraints=VideoTransportConstraints(
                pending_float_consumers=frozenset(
                    FloatVideoConsumer(value) for value in video_payload.get("pending_float_consumers", [])
                )
            ),
            schema_version=1,
        ),
        prepared_for_transport=True,
    )
    media.validate()
    return media


def _pack_diffusion_fields(
    output: DiffusionOutput,
    d2h_stream: torch.Stream | None = None,
    created: list[dict[str, Any]] | None = None,
    pending_updates: list[tuple[DiffusionOutput, str, object]] | None = None,
    *,
    enable_registered_shm: bool = False,
) -> DiffusionOutput:
    if output.media is not None and output.output is not None:
        raise ValueError("DiffusionOutput cannot contain both media and legacy output")
    if output.media is not None and not isinstance(output.media, DiffusionMediaOutput):
        raise TypeError(f"DiffusionOutput.media must be DiffusionMediaOutput, got {type(output.media).__name__}")

    # Pack into locals first so a mid-payload failure leaves ``output`` unmutated
    # (never enqueue a half-packed object); the caller unlinks any SHM segments
    # accumulated in ``created`` on failure.
    packed_output = _UNSET
    if enable_registered_shm and output.video_output_index is not None:
        index = output.video_output_index
        values = output.output
        if not isinstance(values, tuple) or type(index) is not int or not 0 <= index < len(values):
            raise ValueError("video_output_index must identify a video tensor in the legacy output tuple")
        if not isinstance(values[index], torch.Tensor):
            raise TypeError("The legacy video output must be a tensor")
        packed_output = tuple(
            _pack_tensor_if_large(value, d2h_stream=d2h_stream, created=created, enable_registered_shm=True)
            if position == index
            else _pack_value_if_large(value, d2h_stream=d2h_stream, created=created)
            for position, value in enumerate(values)
        )
    elif output.output is not None:
        packed_output = _pack_value_if_large(output.output, d2h_stream=d2h_stream, created=created)
    packed_media = _UNSET
    if output.media is not None:
        packed_media = _pack_diffusion_media(
            output.media, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm
        )
    packed_latents = _UNSET
    if isinstance(output.trajectory_latents, torch.Tensor):
        packed_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream, created=created)
    packed_timesteps = _UNSET
    if isinstance(output.trajectory_timesteps, torch.Tensor):
        packed_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream, created=created)
    packed_log_probs = _UNSET
    if isinstance(output.trajectory_log_probs, torch.Tensor):
        packed_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream, created=created)

    updates = [
        (output, field_name, value)
        for field_name, value in (
            ("output", packed_output),
            ("media", packed_media),
            ("trajectory_latents", packed_latents),
            ("trajectory_timesteps", packed_timesteps),
            ("trajectory_log_probs", packed_log_probs),
        )
        if value is not _UNSET
    ]
    if pending_updates is None:
        for target, field_name, value in updates:
            setattr(target, field_name, value)
    else:
        pending_updates.extend(updates)
    return output


def _is_rpc_result_envelope(output: object) -> TypeGuard[dict[str, Any]]:
    return isinstance(output, dict) and output.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE


def pack_diffusion_output_shm(
    output: object,
    d2h_stream: torch.Stream | None = None,
    *,
    enable_registered_shm: bool = False,
) -> object:
    """Replace large tensors in diffusion worker outputs with SHM handles.

    Supports a bare ``DiffusionOutput``, a wrapper object carrying one in
    ``.result`` (for example ``RunnerOutput``), an RPC result envelope carrying
    the diffusion output in ``["result"]``, a batch wrapper carrying
    ``RunnerOutput`` objects in ``.runner_outputs``, or a DP-tagged dict
    ``{"dp_rank": int, "output": DiffusionOutput}`` used by DP multi-concurrency.

    If *d2h_stream* is provided, D2H copies use that stream (non-blocking on
    the default stream).  The caller must synchronize *d2h_stream* afterward.

    ``enable_registered_shm`` opts explicitly identified video tensors into
    direct CUDA-to-SHM transfer and borrowed CPU views. Other payloads keep
    their existing transport. Registration is per transfer, without pooling.

    Packing is failure-atomic: every SHM segment created during the call is
    tracked and unlinked if any part of the payload fails to pack, and field
    mutations are committed only after the complete payload succeeds.
    """
    created: list[dict[str, Any]] = []
    pending_updates: list[tuple[DiffusionOutput, str, object]] = []
    try:
        packed = _pack_output_shm(
            output,
            d2h_stream=d2h_stream,
            created=created,
            pending_updates=pending_updates,
            enable_registered_shm=enable_registered_shm,
        )
    except BaseException:
        _unlink_shm_handles(created)
        raise

    # Commit mutations only after every field and batch entry packed. Until
    # this point, a failure leaves the complete caller-owned payload untouched.
    for target, field_name, value in pending_updates:
        setattr(target, field_name, value)
    return packed


def _iter_diffusion_outputs(output: object) -> Iterator[DiffusionOutput]:
    """Yield every packable ``DiffusionOutput`` leaf in a worker payload."""
    if isinstance(output, DiffusionOutput):
        yield output
        return

    # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput}
    if isinstance(output, dict) and "dp_rank" in output and "output" in output:
        inner = output["output"]
        if isinstance(inner, DiffusionOutput):
            yield inner
        return

    if _is_rpc_result_envelope(output):
        yield from _iter_diffusion_outputs(output.get("result"))
        return

    result = getattr(output, "result", None)
    if isinstance(result, DiffusionOutput):
        yield result

    runner_outputs = getattr(output, "runner_outputs", None)
    if isinstance(runner_outputs, list):
        for runner_output in runner_outputs:
            yield from _iter_diffusion_outputs(runner_output)


def payload_carries_typed_media(output: object) -> bool:
    """True if *output* carries any typed ``DiffusionOutput.media`` payload."""
    return any(item.media is not None for item in _iter_diffusion_outputs(output))


def _pack_output_shm(
    output: object,
    *,
    d2h_stream: torch.Stream | None,
    created: list[dict[str, Any]],
    pending_updates: list[tuple[DiffusionOutput, str, object]],
    enable_registered_shm: bool = False,
) -> object:
    transport_options = {"enable_registered_shm": True} if enable_registered_shm else {}
    for item in _iter_diffusion_outputs(output):
        _pack_diffusion_fields(
            item,
            d2h_stream=d2h_stream,
            created=created,
            pending_updates=pending_updates,
            **transport_options,
        )
    return output


def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
    output.output = _unpack_if_shm_handle(output.output)
    if isinstance(output.media, dict):
        output.media = _unpack_diffusion_media(output.media)
    output.trajectory_latents = _unpack_if_shm_handle(output.trajectory_latents)
    output.trajectory_timesteps = _unpack_if_shm_handle(output.trajectory_timesteps)
    output.trajectory_log_probs = _unpack_if_shm_handle(output.trajectory_log_probs)
    return output


def unpack_diffusion_output_shm(output: object) -> object:
    """Reconstruct tensors from SHM handles in diffusion worker outputs."""
    if isinstance(output, DiffusionOutput):
        return _unpack_diffusion_fields(output)

    # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput}
    if isinstance(output, dict) and "dp_rank" in output and "output" in output:
        inner = output["output"]
        if isinstance(inner, DiffusionOutput):
            output["output"] = _unpack_diffusion_fields(inner)
        return output

    if _is_rpc_result_envelope(output):
        result = output.get("result")
        output["result"] = unpack_diffusion_output_shm(result)
        return output

    result = getattr(output, "result", None)
    if isinstance(result, DiffusionOutput):
        setattr(output, "result", _unpack_diffusion_fields(result))

    runner_outputs = getattr(output, "runner_outputs", None)
    if isinstance(runner_outputs, list):
        for runner_output in runner_outputs:
            unpack_diffusion_output_shm(runner_output)
    return output