"""CUDA-Enabled Shared Memory Ring Buffer v2.0 - Hybrid CPU/GPU Support.

This module implements a hybrid ring buffer supporting BOTH:
1. CPU → SHM → GPU path (via cudaHostRegister, ~15 GB/s PCIe)
2. GPU → GPU path (via CUDA IPC, ~720 GB/s HBM3 on H100)

The buffer auto-selects the optimal path based on data source while
maintaining full cross-container compatibility.

================================================================================
WHAT IS CUDA IPC AND HOW IT WORKS
================================================================================

CUDA IPC (Inter-Process Communication) allows GPU memory to be shared between
different processes WITHOUT copying through CPU. Here's how it works:

1. PRODUCER allocates GPU memory and gets an IPC handle:
   ┌─────────────────────────────────────────────────────────────────────────┐
   │  gpu_buffer = cp.empty(size, dtype=cp.uint8)     # Allocate on GPU      │
   │  ipc_handle = cudaIpcGetMemHandle(gpu_buffer)    # Get 64-byte handle   │
   │  store_handle_in_shm(ipc_handle)                 # Share via SHM        │
   └─────────────────────────────────────────────────────────────────────────┘

2. CONSUMER opens the handle and accesses the SAME GPU memory:
   ┌─────────────────────────────────────────────────────────────────────────┐
   │  ipc_handle = read_handle_from_shm()             # Read from SHM        │
   │  gpu_ptr = cudaIpcOpenMemHandle(ipc_handle)      # Map into this process│
   │  # Now gpu_ptr points to the SAME memory as producer!                   │
   └─────────────────────────────────────────────────────────────────────────┘

WHY IS CUDA IPC SO FAST?

  WITHOUT CUDA IPC (through CPU/SHM):
  ┌─────────┐     PCIe      ┌─────────┐     PCIe      ┌─────────┐
  │ GPU A   │ ────────────> │   CPU   │ ────────────> │ GPU B   │
  │ Producer│    ~15 GB/s   │   RAM   │    ~15 GB/s   │ Consumer│
  └─────────┘               └─────────┘               └─────────┘
                     Total: ~7.5 GB/s (2x PCIe transfers)

  WITH CUDA IPC (direct GPU memory):
  ┌─────────────────────────────────────────────────────────────┐
  │                     GPU Memory (HBM3)                        │
  │  ┌─────────────────────────────────────────────────────┐    │
  │  │         Shared Buffer (via IPC handle)              │    │
  │  │                                                     │    │
  │  │   Producer writes ──────> Consumer reads            │    │
  │  │                    ~720 GB/s                        │    │
  │  └─────────────────────────────────────────────────────┘    │
  └─────────────────────────────────────────────────────────────┘
                     Total: ~720 GB/s (HBM bandwidth)

================================================================================
ARCHITECTURE
================================================================================

┌─────────────────────────────────────────────────────────────────────────────┐
│                         CudaShmRingBuffer v2.0                              │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  ┌─────────────────────────────────────────────────────────────────────┐    │
│  │ POSIX SHM (Always Present) - Cross-Container Metadata & CPU Path    │    │
│  │  • Header: write_idx, dimensions, format, slot_count                │    │
│  │  • Per-slot metadata: frame_idx, seq_start, seq_end, data_location  │    │
│  │  • Frame data slots (for CPU path)                                  │    │
│  │  • CUDA IPC handles (for GPU path)                                  │    │
│  └─────────────────────────────────────────────────────────────────────┘    │
│                              │                                              │
│              ┌───────────────┴───────────────┐                              │
│              ▼                               ▼                              │
│  ┌─────────────────────────┐    ┌─────────────────────────────┐            │
│  │ CPU Path (~15 GB/s)     │    │ GPU Path (~720 GB/s)        │            │
│  │                         │    │                             │            │
│  │ write_frame(np.array)   │    │ write_frame_gpu(cp.array)   │            │
│  │      ↓ memcpy           │    │      ↓ GPU memcpy           │            │
│  │ SHM (cudaHostRegister)  │    │ GPU Buffer (CUDA IPC)       │            │
│  │      ↓ PCIe DMA         │    │      ↓ HBM bandwidth        │            │
│  │ get_frame_gpu_tensor()  │    │ get_frame_gpu_tensor()      │            │
│  └─────────────────────────┘    └─────────────────────────────┘            │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

================================================================================
CROSS-CONTAINER GPU SHARING (CUDA IPC)
================================================================================

┌─────────────────────────────────────────────────────────────────────────────┐
│ Container A (Producer)              Container B (Consumer)                  │
│ ┌─────────────────────────┐         ┌─────────────────────────┐            │
│ │ 1. Alloc GPU buffer     │         │ 1. Read IPC handle from │            │
│ │ 2. Get IPC handle       │─────────│    SHM                  │            │
│ │ 3. Store handle in SHM  │  SHM    │ 2. ipcOpenMemHandle()   │            │
│ │ 4. Write to GPU buffer  │         │ 3. Access SAME GPU mem  │            │
│ └─────────────────────────┘         └─────────────────────────┘            │
│                                                                             │
│ DOCKER REQUIREMENTS:                                                        │
│   • ipc: "host"           # Share /dev/shm for IPC handles                 │
│   • Same physical GPU     # Both containers see same GPU                   │
│   • NVIDIA driver >= 450  # CUDA IPC support                               │
└─────────────────────────────────────────────────────────────────────────────┘

================================================================================
USAGE EXAMPLES
================================================================================

EXAMPLE 1: CPU Producer → GPU Consumer (Video Decoder → Inference)
──────────────────────────────────────────────────────────────────
    # Producer (video decoder, runs on CPU)
    buffer = CudaShmRingBuffer(
        camera_id="cam_001",
        width=1920, height=1080,
        create=True,
        enable_gpu_ipc=False  # CPU-only producer
    )

    while True:
        frame = decode_video_frame()  # numpy array
        buffer.write_frame(frame)     # ~18 GB/s to SHM

    # Consumer (inference server, runs on GPU)
    buffer = CudaShmRingBuffer(
        camera_id="cam_001",
        width=1920, height=1080,
        create=False
    )

    while True:
        frame_idx = buffer.get_current_frame_idx()
        gpu_tensor = buffer.get_frame_gpu_tensor(frame_idx)  # ~15 GB/s from SHM
        result = model(gpu_tensor)

EXAMPLE 2: GPU Producer → GPU Consumer (Maximum Performance)
────────────────────────────────────────────────────────────
    # Producer (GPU preprocessing)
    buffer = CudaShmRingBuffer(
        camera_id="cam_001",
        width=1920, height=1080,
        create=True,
        enable_gpu_ipc=True  # Enable GPU→GPU path
    )

    while True:
        gpu_frame = preprocess_on_gpu()      # CuPy array
        buffer.write_frame_gpu(gpu_frame)    # ~720 GB/s via CUDA IPC!

    # Consumer (inference server, SAME or DIFFERENT container)
    buffer = CudaShmRingBuffer(
        camera_id="cam_001",
        width=1920, height=1080,
        create=False
    )

    while True:
        frame_idx = buffer.get_current_frame_idx()
        gpu_tensor = buffer.get_frame_gpu_tensor(frame_idx)  # ~720 GB/s!
        result = model(gpu_tensor)

================================================================================
WHEN TO USE EACH PATH
================================================================================

┌─────────────────────────────────────────────────────────────────────────────┐
│ USE CPU PATH (write_frame) WHEN:                                            │
│   ✓ Producer doesn't have GPU (video decoder, camera driver, CPU app)      │
│   ✓ Consumer may not have GPU                                              │
│   ✓ Cross-machine compatibility needed                                     │
│   ✓ Simpler deployment (no CUDA IPC setup)                                 │
│                                                                             │
│ USE GPU IPC PATH (write_frame_gpu) WHEN:                                    │
│   ✓ Both producer and consumer have GPU access                             │
│   ✓ Running on same physical machine                                       │
│   ✓ Need maximum throughput (46x faster than PCIe!)                        │
│   ✓ Latency-critical applications (16us vs 400us)                          │
└─────────────────────────────────────────────────────────────────────────────┘

Version History:
- v1.0: Initial CPU → SHM → GPU implementation
- v1.1: Added health check APIs, multi-GPU support
- v2.0: Added GPU → GPU via CUDA IPC, hybrid path auto-selection
"""

__version__ = "2.0.0"

import logging
import os
import struct
import time
from multiprocessing import shared_memory
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import numpy as np

if TYPE_CHECKING:
    import cupy as cp
    import torch


class CudaRingBufferError(Exception):
    """Base exception for CUDA ring buffer errors."""
    pass


class CudaNotAvailableError(CudaRingBufferError):
    """Raised when CUDA is not available."""
    pass


class CudaIpcError(CudaRingBufferError):
    """Raised when CUDA IPC operations fail."""
    pass


# Data location flags for per-slot metadata
DATA_LOCATION_SHM = 0      # Frame data in SHM (CPU path)
DATA_LOCATION_GPU_IPC = 1  # Frame data in GPU via IPC


class CudaShmRingBuffer:
    """Hybrid CUDA ring buffer with automatic CPU/GPU path selection.

    This buffer supports three modes of operation:

    1. CPU Producer → GPU Consumer (Original):
       - Producer calls write_frame(numpy_array)
       - Data stored in CUDA-registered SHM
       - Consumer reads via PCIe DMA (~14 GB/s)

    2. GPU Producer → GPU Consumer via SHM:
       - Producer calls write_frame_gpu(cupy_array, use_ipc=False)
       - Data copied GPU→Host→SHM via pinned memory
       - Consumer reads via PCIe DMA (~14 GB/s)

    3. GPU Producer → GPU Consumer via CUDA IPC (Maximum Performance):
       - Producer calls write_frame_gpu(cupy_array, use_ipc=True)
       - Data stays in GPU memory, IPC handle shared via SHM
       - Consumer accesses SAME GPU memory (~700 GB/s)

    The consumer automatically detects which path was used and reads accordingly.

    Example (Producer with auto-selection):
        buffer = CudaShmRingBuffer(
            camera_id="cam_001",
            width=1920, height=1080,
            create=True,
            enable_gpu_ipc=True  # Enable GPU→GPU path
        )

        # CPU data → uses SHM path automatically
        buffer.write_frame(numpy_frame)

        # GPU data → uses CUDA IPC path automatically
        buffer.write_frame_gpu(cupy_frame)

    Example (Consumer - automatic path detection):
        buffer = CudaShmRingBuffer(
            camera_id="cam_001",
            width=1920, height=1080,
            create=False
        )

        # Auto-detects SHM or IPC path based on where producer stored data
        gpu_tensor = buffer.get_frame_gpu_tensor(frame_idx)
    """

    # Header layout (extended for v2.0)
    # Format: write_idx(Q), width(I), height(I), format(I), slot_count(I),
    #         last_ts_ns(Q), version(I), flags(I), gpu_ipc_enabled(I), reserved(I)
    HEADER_FORMAT = '<QIIIIQIIII'
    HEADER_SIZE = 64

    # Per-slot metadata (extended for v2.0)
    # Format: frame_idx(Q), seq_start(I), seq_end(I), data_location(I), reserved(I)
    SLOT_METADATA_FORMAT = '<QIIII'
    SLOT_METADATA_SIZE = 24

    # CUDA IPC handle size (64 bytes for cudaIpcMemHandle_t)
    CUDA_IPC_HANDLE_SIZE = 64

    PAGE_SIZE = 4096
    VERSION = 2  # Buffer format version

    # Frame format constants
    FORMAT_NV12 = 0
    FORMAT_RGB = 1
    FORMAT_BGR = 2

    FORMAT_NAMES = {
        FORMAT_NV12: "NV12",
        FORMAT_RGB: "RGB",
        FORMAT_BGR: "BGR",
    }

    # CUDA host register flags
    CUDA_HOST_REGISTER_DEFAULT = 0x00
    CUDA_HOST_REGISTER_PORTABLE = 0x01
    CUDA_HOST_REGISTER_MAPPED = 0x02

    def __init__(
        self,
        camera_id: str,
        width: int,
        height: int,
        frame_format: int = FORMAT_BGR,
        slot_count: int = 16,
        create: bool = True,
        use_zero_copy: bool = True,
        cuda_device: int = 0,
        fallback_to_cpu: bool = True,
        enable_gpu_ipc: bool = True,
        shm_name: Optional[str] = None,
    ):
        """Initialize hybrid CUDA ring buffer.

        Args:
            camera_id: Unique camera/stream identifier
            width: Frame width in pixels
            height: Frame height in pixels
            frame_format: One of FORMAT_NV12, FORMAT_RGB, FORMAT_BGR
            slot_count: Number of frame slots (8-16 recommended for GPU)
            create: True for producer (creates buffer), False for consumer
            use_zero_copy: Enable zero-copy GPU access via cudaHostGetDevicePointer
            cuda_device: CUDA device ID to use
            fallback_to_cpu: Fall back to CPU-only mode if CUDA unavailable
            enable_gpu_ipc: Enable GPU→GPU path via CUDA IPC (requires same GPU access)
            shm_name: Direct SHM name (optional, auto-generated if not provided)
        """
        self.logger = logging.getLogger(f"{__name__}.CudaShmRingBuffer")

        # Validate inputs
        if frame_format not in self.FORMAT_NAMES:
            raise ValueError(f"Invalid frame_format: {frame_format}")
        if width <= 0 or height <= 0:
            raise ValueError(f"Invalid dimensions: {width}x{height}")
        if slot_count < 2:
            raise ValueError(f"slot_count must be >= 2, got {slot_count}")

        self.camera_id = camera_id
        self.width = width
        self.height = height
        self.frame_format = frame_format
        self.slot_count = slot_count
        self._is_producer = create
        self._use_zero_copy = use_zero_copy
        self._cuda_device = cuda_device
        self._fallback_to_cpu = fallback_to_cpu
        self._enable_gpu_ipc = enable_gpu_ipc

        # Calculate sizes
        self.frame_size = self._calculate_frame_size(width, height, frame_format)
        self._aligned_slot_size = self._align_to_page(self.frame_size)
        self._slot_metadata_total = self.SLOT_METADATA_SIZE * slot_count
        self._ipc_handles_offset = self.HEADER_SIZE + self._slot_metadata_total
        self._ipc_handles_size = self.CUDA_IPC_HANDLE_SIZE * slot_count
        self._frame_data_offset = self._ipc_handles_offset + self._ipc_handles_size
        self._total_shm_size = self._frame_data_offset + (self._aligned_slot_size * slot_count)

        # Generate SHM name
        self.shm_name = shm_name if shm_name else self._generate_shm_name(camera_id)

        # Producer-side cached counters
        self._cached_write_idx: int = 0
        self._cached_slot_seq: List[int] = [0] * slot_count

        # Shared memory state
        self._shm: Optional[shared_memory.SharedMemory] = None
        self._buffer_array: Optional[np.ndarray] = None

        # CUDA state
        self._cuda_available = False
        self._cupy_available = False
        self._torch_available = False
        self._cuda_registered = False
        self._device_ptr: Optional[int] = None
        self._cp = None
        self._torch = None

        # GPU IPC state (for GPU→GPU path)
        self._gpu_ipc_available = False
        self._gpu_buffer: Optional[Any] = None  # CuPy array for GPU-side ring buffer
        self._gpu_buffer_ptr: int = 0
        self._consumer_ipc_ptrs: Dict[int, int] = {}  # slot -> opened IPC pointer

        # Initialize
        self._init_shared_memory(create)
        self._init_cuda_registration()
        if enable_gpu_ipc:
            self._init_gpu_ipc(create)

        self.logger.info(
            f"CudaShmRingBuffer v{__version__} {'created' if create else 'attached'}: "
            f"name={self.shm_name}, {width}x{height} {self.FORMAT_NAMES[frame_format]}, "
            f"{slot_count} slots, cuda={self._cuda_available}, "
            f"zero_copy={self._device_ptr is not None}, "
            f"gpu_ipc={self._gpu_ipc_available}"
        )

    def _generate_shm_name(self, camera_id: str) -> str:
        """Generate SHM name from camera ID."""
        safe_id = camera_id.replace("/", "_").replace("\\", "_").replace(":", "_")
        safe_id = "".join(c for c in safe_id if c.isalnum() or c == "_")
        return f"cuda_shm_v2_{safe_id[:160]}"

    def _calculate_frame_size(self, width: int, height: int, frame_format: int) -> int:
        """Calculate frame size in bytes."""
        if frame_format == self.FORMAT_NV12:
            return int(width * height * 1.5)
        elif frame_format in (self.FORMAT_RGB, self.FORMAT_BGR):
            return width * height * 3
        raise ValueError(f"Unknown frame format: {frame_format}")

    def _align_to_page(self, size: int) -> int:
        """Align size to page boundary."""
        return ((size + self.PAGE_SIZE - 1) // self.PAGE_SIZE) * self.PAGE_SIZE

    # =========================================================================
    # Shared Memory Initialization
    # =========================================================================

    def _init_shared_memory(self, create: bool) -> None:
        """Initialize POSIX shared memory."""
        if create:
            try:
                self._shm = shared_memory.SharedMemory(
                    name=self.shm_name, create=True, size=self._total_shm_size
                )
            except FileExistsError:
                self.logger.warning(f"SHM {self.shm_name} exists, recreating")
                try:
                    old_shm = shared_memory.SharedMemory(name=self.shm_name, create=False)
                    old_shm.close()
                    old_shm.unlink()
                except Exception as e:
                    self.logger.warning(f"Could not unlink old SHM: {e}")
                self._shm = shared_memory.SharedMemory(
                    name=self.shm_name, create=True, size=self._total_shm_size
                )

            self._buffer_array = np.ndarray(
                shape=(self._total_shm_size,), dtype=np.uint8, buffer=self._shm.buf
            )

            # Initialize header
            self._write_header(
                write_idx=0, width=self.width, height=self.height,
                frame_format=self.frame_format, slot_count=self.slot_count,
                last_ts_ns=int(time.time() * 1e9), version=self.VERSION,
                flags=0, gpu_ipc_enabled=1 if self._enable_gpu_ipc else 0
            )

            # Initialize slot metadata
            for slot in range(self.slot_count):
                self._write_slot_metadata(slot, frame_idx=0, seq_start=0, seq_end=0,
                                          data_location=DATA_LOCATION_SHM)
        else:
            try:
                self._shm = shared_memory.SharedMemory(name=self.shm_name, create=False)
                self._buffer_array = np.ndarray(
                    shape=(self._shm.size,), dtype=np.uint8, buffer=self._shm.buf
                )

                # Validate header
                header = self._read_header()
                if header['width'] != self.width or header['height'] != self.height:
                    self.logger.warning(
                        f"Dimension mismatch: expected {self.width}x{self.height}, "
                        f"got {header['width']}x{header['height']}"
                    )
                self._enable_gpu_ipc = header.get('gpu_ipc_enabled', 0) == 1

            except FileNotFoundError:
                raise FileNotFoundError(
                    f"SHM '{self.shm_name}' not found. Start producer first."
                )

    def _init_cuda_registration(self) -> None:
        """Register SHM with CUDA for pinned memory access."""
        try:
            import cupy as cp
            self._cp = cp
            self._cupy_available = True
            self._cuda_available = True

            with cp.cuda.Device(self._cuda_device):
                shm_ptr = self._buffer_array.ctypes.data
                flags = self.CUDA_HOST_REGISTER_PORTABLE
                if self._use_zero_copy:
                    flags |= self.CUDA_HOST_REGISTER_MAPPED

                try:
                    cp.cuda.runtime.hostRegister(shm_ptr, self._total_shm_size, flags)
                    self._cuda_registered = True
                except cp.cuda.runtime.CUDARuntimeError as e:
                    if "already registered" in str(e).lower() or "712" in str(e):
                        self._cuda_registered = True
                    else:
                        raise

                if self._use_zero_copy and self._cuda_registered:
                    try:
                        self._device_ptr = cp.cuda.runtime.hostGetDevicePointer(shm_ptr, 0)
                    except Exception as e:
                        self.logger.warning(f"cudaHostGetDevicePointer failed: {e}")

            self.logger.info(
                f"CUDA registered: device={self._cuda_device}, "
                f"registered={self._cuda_registered}, zero_copy={self._device_ptr is not None}"
            )
            return

        except ImportError:
            self.logger.debug("CuPy not available")
        except Exception as e:
            self.logger.warning(f"CUDA registration failed: {e}")

        # Try PyTorch fallback
        try:
            import torch
            if torch.cuda.is_available():
                self._torch = torch
                self._torch_available = True
                self._cuda_available = True
                self.logger.info("CUDA via PyTorch (limited functionality)")
                return
        except ImportError:
            pass

        if not self._cuda_available and not self._fallback_to_cpu:
            raise CudaNotAvailableError("CUDA not available and fallback_to_cpu=False")

    def _init_gpu_ipc(self, create: bool) -> None:
        """Initialize GPU buffer with CUDA IPC for GPU→GPU path."""
        if not self._cupy_available:
            self.logger.debug("GPU IPC requires CuPy")
            return

        cp = self._cp

        try:
            with cp.cuda.Device(self._cuda_device):
                if create:
                    # Producer: Allocate GPU ring buffer
                    total_gpu_size = self._aligned_slot_size * self.slot_count
                    self._gpu_buffer = cp.empty(total_gpu_size, dtype=cp.uint8)
                    self._gpu_buffer_ptr = self._gpu_buffer.data.ptr

                    # Get IPC handle for each slot and store in SHM
                    for slot in range(self.slot_count):
                        slot_ptr = self._gpu_buffer_ptr + (slot * self._aligned_slot_size)
                        try:
                            # Note: IPC handle is for the base allocation, offset calculated at read
                            ipc_handle = cp.cuda.runtime.ipcGetMemHandle(self._gpu_buffer_ptr)
                            self._write_ipc_handle(slot, bytes(ipc_handle))
                        except Exception as e:
                            self.logger.warning(f"Failed to get IPC handle for slot {slot}: {e}")
                            return

                    self._gpu_ipc_available = True
                    self.logger.info(
                        f"GPU IPC initialized: {total_gpu_size / 1024 / 1024:.1f} MB GPU buffer"
                    )
                else:
                    # Consumer: Will open IPC handles on-demand when reading
                    self._gpu_ipc_available = self._enable_gpu_ipc
                    if self._gpu_ipc_available:
                        self.logger.info("GPU IPC consumer mode ready")

        except Exception as e:
            self.logger.warning(f"GPU IPC initialization failed: {e}")
            self._gpu_ipc_available = False

    # =========================================================================
    # Header and Metadata Operations
    # =========================================================================

    def _write_header(self, write_idx: int, width: int, height: int, frame_format: int,
                      slot_count: int, last_ts_ns: int, version: int, flags: int,
                      gpu_ipc_enabled: int) -> None:
        """Write header to SHM."""
        header_data = struct.pack(
            self.HEADER_FORMAT,
            write_idx, width, height, frame_format, slot_count,
            last_ts_ns, version, flags, gpu_ipc_enabled, 0
        )
        self._buffer_array[:self.HEADER_SIZE] = np.frombuffer(
            header_data.ljust(self.HEADER_SIZE, b'\x00'), dtype=np.uint8
        )

    def _read_header(self) -> Dict:
        """Read header from SHM."""
        header_bytes = bytes(self._buffer_array[:struct.calcsize(self.HEADER_FORMAT)])
        values = struct.unpack(self.HEADER_FORMAT, header_bytes)
        return {
            'write_idx': values[0],
            'width': values[1],
            'height': values[2],
            'format': values[3],
            'slot_count': values[4],
            'last_ts_ns': values[5],
            'version': values[6],
            'flags': values[7],
            'gpu_ipc_enabled': values[8],
        }

    def _update_write_idx(self, new_idx: int) -> None:
        """Update write_idx in header."""
        self._buffer_array[:8] = np.frombuffer(struct.pack('<Q', new_idx), dtype=np.uint8)

    def _update_last_ts_ns(self, ts_ns: int) -> None:
        """Update heartbeat timestamp."""
        offset = struct.calcsize('<QIIII')  # After write_idx, width, height, format, slot_count
        self._buffer_array[offset:offset + 8] = np.frombuffer(
            struct.pack('<Q', ts_ns), dtype=np.uint8
        )

    def _get_slot_metadata_offset(self, slot: int) -> int:
        """Get byte offset for slot metadata."""
        return self.HEADER_SIZE + (slot * self.SLOT_METADATA_SIZE)

    def _write_slot_metadata(self, slot: int, frame_idx: int, seq_start: int,
                             seq_end: int, data_location: int) -> None:
        """Write complete slot metadata."""
        offset = self._get_slot_metadata_offset(slot)
        data = struct.pack(self.SLOT_METADATA_FORMAT, frame_idx, seq_start, seq_end,
                           data_location, 0)
        self._buffer_array[offset:offset + self.SLOT_METADATA_SIZE] = np.frombuffer(
            data, dtype=np.uint8
        )

    def _read_slot_metadata(self, slot: int) -> Dict:
        """Read slot metadata."""
        offset = self._get_slot_metadata_offset(slot)
        data = bytes(self._buffer_array[offset:offset + self.SLOT_METADATA_SIZE])
        values = struct.unpack(self.SLOT_METADATA_FORMAT, data)
        return {
            'frame_idx': values[0],
            'seq_start': values[1],
            'seq_end': values[2],
            'data_location': values[3],
        }

    def _write_slot_seq_start(self, slot: int, seq: int) -> None:
        """Write seq_start for torn frame detection."""
        offset = self._get_slot_metadata_offset(slot) + 8
        self._buffer_array[offset:offset + 4] = np.frombuffer(
            struct.pack('<I', seq & 0xFFFFFFFF), dtype=np.uint8
        )

    def _write_slot_seq_end(self, slot: int, seq: int) -> None:
        """Write seq_end for torn frame detection."""
        offset = self._get_slot_metadata_offset(slot) + 12
        self._buffer_array[offset:offset + 4] = np.frombuffer(
            struct.pack('<I', seq & 0xFFFFFFFF), dtype=np.uint8
        )

    def _write_slot_data_location(self, slot: int, location: int) -> None:
        """Write data location flag."""
        offset = self._get_slot_metadata_offset(slot) + 16
        self._buffer_array[offset:offset + 4] = np.frombuffer(
            struct.pack('<I', location), dtype=np.uint8
        )

    def _write_slot_frame_idx(self, slot: int, frame_idx: int) -> None:
        """Write frame_idx to slot."""
        offset = self._get_slot_metadata_offset(slot)
        self._buffer_array[offset:offset + 8] = np.frombuffer(
            struct.pack('<Q', frame_idx), dtype=np.uint8
        )

    def _read_slot_frame_idx(self, slot: int) -> int:
        """Read frame_idx from slot."""
        offset = self._get_slot_metadata_offset(slot)
        return struct.unpack('<Q', bytes(self._buffer_array[offset:offset + 8]))[0]

    def _read_slot_seq_start(self, slot: int) -> int:
        """Read seq_start."""
        offset = self._get_slot_metadata_offset(slot) + 8
        return struct.unpack('<I', bytes(self._buffer_array[offset:offset + 4]))[0]

    def _read_slot_seq_end(self, slot: int) -> int:
        """Read seq_end."""
        offset = self._get_slot_metadata_offset(slot) + 12
        return struct.unpack('<I', bytes(self._buffer_array[offset:offset + 4]))[0]

    def _read_slot_data_location(self, slot: int) -> int:
        """Read data location flag."""
        offset = self._get_slot_metadata_offset(slot) + 16
        return struct.unpack('<I', bytes(self._buffer_array[offset:offset + 4]))[0]

    def _write_ipc_handle(self, slot: int, handle_bytes: bytes) -> None:
        """Write CUDA IPC handle for slot."""
        offset = self._ipc_handles_offset + (slot * self.CUDA_IPC_HANDLE_SIZE)
        handle_data = handle_bytes[:self.CUDA_IPC_HANDLE_SIZE].ljust(
            self.CUDA_IPC_HANDLE_SIZE, b'\x00'
        )
        self._buffer_array[offset:offset + self.CUDA_IPC_HANDLE_SIZE] = np.frombuffer(
            handle_data, dtype=np.uint8
        )

    def _read_ipc_handle(self, slot: int) -> bytes:
        """Read CUDA IPC handle for slot."""
        offset = self._ipc_handles_offset + (slot * self.CUDA_IPC_HANDLE_SIZE)
        return bytes(self._buffer_array[offset:offset + self.CUDA_IPC_HANDLE_SIZE])

    def _get_frame_offset(self, slot: int) -> int:
        """Get byte offset for frame data in SHM."""
        return self._frame_data_offset + (slot * self._aligned_slot_size)

    # =========================================================================
    # Producer Operations - CPU Path
    # =========================================================================

    def write_frame(self, raw_bytes: Union[bytes, memoryview, np.ndarray]) -> Tuple[int, int]:
        """Write frame from CPU to SHM (original path).

        Args:
            raw_bytes: Frame data as bytes, memoryview, or numpy array

        Returns:
            Tuple of (frame_idx, slot_idx)
        """
        if not self._is_producer:
            raise RuntimeError("write_frame() only callable on producer")

        # Handle input types
        if isinstance(raw_bytes, np.ndarray):
            if raw_bytes.flags['C_CONTIGUOUS']:
                raw_data = raw_bytes.ravel().view(np.uint8)
            else:
                raw_data = np.ascontiguousarray(raw_bytes).ravel().view(np.uint8)
        elif isinstance(raw_bytes, memoryview):
            raw_data = np.frombuffer(raw_bytes, dtype=np.uint8)
        else:
            raw_data = np.frombuffer(raw_bytes, dtype=np.uint8)

        if len(raw_data) != self.frame_size:
            raise ValueError(f"Frame size mismatch: expected {self.frame_size}, got {len(raw_data)}")

        # Get next slot
        self._cached_write_idx += 1
        frame_idx = self._cached_write_idx
        slot = frame_idx % self.slot_count

        # Start write (odd sequence)
        self._cached_slot_seq[slot] += 1
        self._write_slot_seq_start(slot, self._cached_slot_seq[slot])
        self._write_slot_data_location(slot, DATA_LOCATION_SHM)

        # Write frame data to SHM
        frame_offset = self._get_frame_offset(slot)
        self._buffer_array[frame_offset:frame_offset + self.frame_size] = raw_data

        # Complete write (even sequence)
        self._write_slot_frame_idx(slot, frame_idx)
        self._cached_slot_seq[slot] += 1
        self._write_slot_seq_end(slot, self._cached_slot_seq[slot])
        self._write_slot_seq_start(slot, self._cached_slot_seq[slot])

        # Memory barrier
        _ = bytes(self._shm.buf[0:1])

        # Update header
        self._update_write_idx(frame_idx)
        self._update_last_ts_ns(int(time.time() * 1e9))

        return frame_idx, slot

    # =========================================================================
    # Producer Operations - GPU Path
    # =========================================================================

    def write_frame_gpu(
        self,
        gpu_array: Any,
        use_ipc: bool = True,
        stream: Optional[Any] = None
    ) -> Tuple[int, int]:
        """Write frame from GPU memory.

        This method automatically selects the optimal path:
        - use_ipc=True (default): GPU→GPU via CUDA IPC (~700 GB/s)
        - use_ipc=False: GPU→SHM via pinned memory (~14 GB/s)

        Args:
            gpu_array: CuPy ndarray on GPU
            use_ipc: Use CUDA IPC path if available (recommended)
            stream: CUDA stream for async operations

        Returns:
            Tuple of (frame_idx, slot_idx)
        """
        if not self._is_producer:
            raise RuntimeError("write_frame_gpu() only callable on producer")

        if not self._cupy_available:
            raise RuntimeError("CuPy required for GPU writes")

        cp = self._cp

        # Flatten and validate
        flat_array = gpu_array.ravel().view(cp.uint8)
        if flat_array.size != self.frame_size:
            raise ValueError(f"Frame size mismatch: expected {self.frame_size}, got {flat_array.size}")

        # Get next slot
        self._cached_write_idx += 1
        frame_idx = self._cached_write_idx
        slot = frame_idx % self.slot_count

        # Start write (odd sequence)
        self._cached_slot_seq[slot] += 1
        self._write_slot_seq_start(slot, self._cached_slot_seq[slot])

        with cp.cuda.Device(self._cuda_device):
            if use_ipc and self._gpu_ipc_available and self._gpu_buffer is not None:
                # GPU → GPU IPC path (maximum performance)
                self._write_slot_data_location(slot, DATA_LOCATION_GPU_IPC)

                # Copy within GPU memory
                slot_offset = slot * self._aligned_slot_size
                dst_slice = self._gpu_buffer[slot_offset:slot_offset + self.frame_size]

                if stream is not None:
                    # Async copy
                    dst_slice.data.copy_from_device_async(
                        flat_array.data, self.frame_size, stream
                    )
                else:
                    cp.copyto(dst_slice, flat_array)
            else:
                # GPU → SHM path (via pinned memory)
                self._write_slot_data_location(slot, DATA_LOCATION_SHM)

                frame_offset = self._get_frame_offset(slot)
                host_slice = self._buffer_array[frame_offset:frame_offset + self.frame_size]

                if stream is not None:
                    flat_array.get(out=host_slice, stream=stream)
                else:
                    host_slice[:] = cp.asnumpy(flat_array)

        # Complete write (even sequence)
        self._write_slot_frame_idx(slot, frame_idx)
        self._cached_slot_seq[slot] += 1
        self._write_slot_seq_end(slot, self._cached_slot_seq[slot])
        self._write_slot_seq_start(slot, self._cached_slot_seq[slot])

        # Memory barrier
        _ = bytes(self._shm.buf[0:1])

        # Update header
        self._update_write_idx(frame_idx)
        self._update_last_ts_ns(int(time.time() * 1e9))

        return frame_idx, slot

    # =========================================================================
    # Consumer Operations - Automatic Path Selection
    # =========================================================================

    def get_frame_gpu_tensor(
        self,
        frame_idx: int,
        copy: bool = True,
        stream: Optional[Any] = None
    ) -> Optional[Any]:
        """Get frame as GPU tensor with automatic path selection.

        Automatically detects whether the frame is in:
        - SHM (CPU path): Reads via PCIe DMA
        - GPU IPC (GPU path): Reads directly from GPU memory

        Args:
            frame_idx: Frame index to read
            copy: For SHM path - copy to device memory (recommended)
            stream: CUDA stream for async operations

        Returns:
            CuPy ndarray on GPU, or None if invalid/torn
        """
        if not self._cupy_available:
            self.logger.warning("CuPy not available")
            return None

        if not self.is_frame_valid(frame_idx):
            return None

        slot = frame_idx % self.slot_count

        # Check for torn frame
        seq_start = self._read_slot_seq_start(slot)
        if seq_start & 1:
            return None

        # Determine data location
        data_location = self._read_slot_data_location(slot)
        cp = self._cp

        with cp.cuda.Device(self._cuda_device):
            if data_location == DATA_LOCATION_GPU_IPC and self._gpu_ipc_available:
                # GPU IPC path - read from GPU memory
                gpu_array = self._read_from_gpu_ipc(slot, stream)
            else:
                # SHM path - read via PCIe DMA
                gpu_array = self._read_from_shm(slot, copy, stream)

        if gpu_array is None:
            return None

        # Verify no torn frame after read
        seq_end = self._read_slot_seq_end(slot)
        if seq_start != seq_end:
            return None

        return gpu_array

    def _read_from_shm(self, slot: int, copy: bool, stream: Optional[Any]) -> Optional[Any]:
        """Read frame from SHM via PCIe DMA."""
        cp = self._cp
        frame_offset = self._get_frame_offset(slot)
        host_array = self._buffer_array[frame_offset:frame_offset + self.frame_size]

        with cp.cuda.Device(self._cuda_device):
            if copy:
                if stream is not None:
                    gpu_array = cp.empty(self.frame_size, dtype=cp.uint8)
                    gpu_array.set(host_array, stream=stream)
                else:
                    gpu_array = cp.asarray(host_array)
            else:
                # Zero-copy via device pointer
                if self._device_ptr is None:
                    self.logger.warning("Zero-copy not available")
                    return None

                mem_ptr = cp.cuda.MemoryPointer(
                    cp.cuda.UnownedMemory(
                        self._device_ptr + frame_offset,
                        self.frame_size,
                        owner=self
                    ),
                    0
                )
                gpu_array = cp.ndarray(
                    shape=(self.frame_size,),
                    dtype=cp.uint8,
                    memptr=mem_ptr
                )

        return gpu_array

    def _read_from_gpu_ipc(self, slot: int, stream: Optional[Any]) -> Optional[Any]:
        """Read frame from GPU memory via CUDA IPC."""
        cp = self._cp

        with cp.cuda.Device(self._cuda_device):
            # Open IPC handle if not already open
            if slot not in self._consumer_ipc_ptrs:
                ipc_handle = self._read_ipc_handle(slot)
                try:
                    # Note: We open the handle to the base GPU buffer
                    ptr = cp.cuda.runtime.ipcOpenMemHandle(
                        ipc_handle,
                        cp.cuda.runtime.cudaIpcMemLazyEnablePeerAccess
                    )
                    self._consumer_ipc_ptrs[slot] = ptr
                except Exception as e:
                    self.logger.warning(f"Failed to open IPC handle for slot {slot}: {e}")
                    return None

            # Calculate offset within the GPU buffer
            base_ptr = self._consumer_ipc_ptrs[slot]
            slot_offset = slot * self._aligned_slot_size

            # Create view into shared GPU memory
            mem = cp.cuda.UnownedMemory(
                base_ptr + slot_offset,
                self.frame_size,
                owner=self
            )
            mem_ptr = cp.cuda.MemoryPointer(mem, 0)

            # Create a copy for safety (the source may be overwritten)
            src_array = cp.ndarray(
                shape=(self.frame_size,),
                dtype=cp.uint8,
                memptr=mem_ptr
            )

            if stream is not None:
                gpu_array = cp.empty(self.frame_size, dtype=cp.uint8)
                gpu_array.data.copy_from_device_async(src_array.data, self.frame_size, stream)
            else:
                gpu_array = src_array.copy()

        return gpu_array

    # =========================================================================
    # Consumer Operations - CPU Path
    # =========================================================================

    def read_frame(self, frame_idx: int) -> Optional[np.ndarray]:
        """Read frame as numpy array (zero-copy view into SHM)."""
        if not self.is_frame_valid(frame_idx):
            return None

        slot = frame_idx % self.slot_count
        data_location = self._read_slot_data_location(slot)

        if data_location == DATA_LOCATION_GPU_IPC:
            # Frame is in GPU memory, need to copy to CPU
            gpu_tensor = self.get_frame_gpu_tensor(frame_idx)
            if gpu_tensor is None:
                return None
            return self._cp.asnumpy(gpu_tensor)

        frame_offset = self._get_frame_offset(slot)
        return self._buffer_array[frame_offset:frame_offset + self.frame_size]

    def read_frame_copy(self, frame_idx: int, max_wait_ms: float = 5.0) -> Optional[bytes]:
        """Read frame with torn frame detection, return safe copy."""
        if not self.is_frame_valid(frame_idx, max_wait_ms=max_wait_ms):
            return None

        slot = frame_idx % self.slot_count
        start_time = time.time()
        max_wait_sec = max_wait_ms / 1000.0

        while True:
            seq_start = self._read_slot_seq_start(slot)
            if seq_start & 1:
                if time.time() - start_time >= max_wait_sec:
                    return None
                time.sleep(0.0002)
                continue

            data_location = self._read_slot_data_location(slot)

            if data_location == DATA_LOCATION_GPU_IPC:
                # Frame in GPU memory
                gpu_tensor = self.get_frame_gpu_tensor(frame_idx)
                if gpu_tensor is None:
                    return None
                frame_data = bytes(self._cp.asnumpy(gpu_tensor))
            else:
                # Frame in SHM
                frame_offset = self._get_frame_offset(slot)
                frame_data = bytes(self._buffer_array[frame_offset:frame_offset + self.frame_size])

            seq_end = self._read_slot_seq_end(slot)
            if seq_start != seq_end:
                if time.time() - start_time >= max_wait_sec:
                    return None
                time.sleep(0.0002)
                continue

            stored_frame_idx = self._read_slot_frame_idx(slot)
            if stored_frame_idx != frame_idx:
                return None

            return frame_data

    # =========================================================================
    # DLPack Support
    # =========================================================================

    def get_frame_dlpack(self, frame_idx: int, copy: bool = True) -> Optional[Any]:
        """Get frame as DLPack capsule for framework-agnostic exchange."""
        if self._cupy_available:
            gpu_array = self.get_frame_gpu_tensor(frame_idx, copy=copy)
            if gpu_array is None:
                return None
            return gpu_array.toDlpack()
        return None

    def get_frame_shaped(
        self,
        frame_idx: int,
        as_gpu: bool = False,
        copy: bool = True
    ) -> Optional[np.ndarray]:
        """Get frame reshaped to (H, W, C)."""
        if as_gpu and self._cupy_available:
            flat_array = self.get_frame_gpu_tensor(frame_idx, copy=copy)
            if flat_array is None:
                return None
        else:
            flat_array = self.read_frame(frame_idx)
            if flat_array is None:
                return None

        if self.frame_format in (self.FORMAT_RGB, self.FORMAT_BGR):
            if as_gpu and self._cupy_available:
                return flat_array.reshape((self.height, self.width, 3))
            return np.asarray(flat_array).reshape((self.height, self.width, 3))
        return flat_array

    # =========================================================================
    # Validation and Status
    # =========================================================================

    def is_frame_valid(self, frame_idx: int, max_wait_ms: float = 5.0) -> bool:
        """Check if frame is valid (not overwritten, not torn)."""
        if frame_idx <= 0:
            return False

        start_time = time.time()
        max_wait_sec = max_wait_ms / 1000.0

        while True:
            header = self._read_header()
            current_write_idx = header['write_idx']

            if current_write_idx - frame_idx >= self.slot_count:
                return False

            if frame_idx > current_write_idx:
                if time.time() - start_time >= max_wait_sec:
                    return False
                time.sleep(0.0001)
                continue

            slot = frame_idx % self.slot_count
            stored_frame_idx = self._read_slot_frame_idx(slot)

            if stored_frame_idx == frame_idx:
                return True

            if frame_idx == current_write_idx:
                if time.time() - start_time >= max_wait_sec:
                    return False
                time.sleep(0.0001)
                continue

            return False

    def get_current_frame_idx(self) -> int:
        """Get latest written frame index."""
        return self._read_header()['write_idx']

    def is_producer_alive(self, timeout_ns: int = 2_000_000_000) -> bool:
        """Check if producer is alive via heartbeat."""
        header = self._read_header()
        return (time.time_ns() - header['last_ts_ns']) < timeout_ns

    def get_health_status(self) -> Dict:
        """Get comprehensive health status."""
        try:
            header = self._read_header()
            return {
                'is_healthy': True,
                'producer_alive': self.is_producer_alive(),
                'frames_written': header['write_idx'],
                'shm_name': self.shm_name,
                'width': self.width,
                'height': self.height,
                'slot_count': self.slot_count,
                'cuda_available': self._cuda_available,
                'cuda_registered': self._cuda_registered,
                'zero_copy_enabled': self._device_ptr is not None,
                'gpu_ipc_available': self._gpu_ipc_available,
                'version': header.get('version', 1),
            }
        except Exception as e:
            return {'is_healthy': False, 'error': str(e)}

    # =========================================================================
    # Properties
    # =========================================================================

    @property
    def is_cuda_available(self) -> bool:
        return self._cuda_available

    @property
    def is_cuda_registered(self) -> bool:
        return self._cuda_registered

    @property
    def is_zero_copy_enabled(self) -> bool:
        return self._device_ptr is not None

    @property
    def is_gpu_ipc_available(self) -> bool:
        return self._gpu_ipc_available

    @property
    def cuda_device(self) -> int:
        return self._cuda_device

    # =========================================================================
    # Cleanup
    # =========================================================================

    def close(self) -> None:
        """Close and free all resources."""
        # Close IPC handles
        if self._consumer_ipc_ptrs and self._cupy_available:
            cp = self._cp
            for slot, ptr in self._consumer_ipc_ptrs.items():
                try:
                    cp.cuda.runtime.ipcCloseMemHandle(ptr)
                except Exception:
                    pass
            self._consumer_ipc_ptrs.clear()

        # Free GPU buffer (producer only)
        if self._is_producer and self._gpu_buffer is not None:
            self._gpu_buffer = None

        # Unregister CUDA host memory
        if self._cuda_registered and self._cupy_available and self._buffer_array is not None:
            try:
                self._cp.cuda.runtime.hostUnregister(self._buffer_array.ctypes.data)
            except Exception:
                pass

        self._cuda_registered = False
        self._device_ptr = None
        self._buffer_array = None

        # Close SHM
        if self._shm is not None:
            try:
                self._shm.close()
                if self._is_producer:
                    try:
                        self._shm.unlink()
                        self.logger.info(f"SHM {self.shm_name} unlinked")
                    except FileNotFoundError:
                        pass
            except Exception as e:
                self.logger.warning(f"Error closing SHM: {e}")
            finally:
                self._shm = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def __repr__(self) -> str:
        return (
            f"CudaShmRingBuffer(name={self.shm_name}, "
            f"{self.width}x{self.height}, slots={self.slot_count}, "
            f"cuda={self._cuda_available}, gpu_ipc={self._gpu_ipc_available})"
        )
