"""Shared Memory Ring Buffer for raw frame storage.

This module implements a lock-free ring buffer using Python's multiprocessing.shared_memory
for storing raw video frames. It enables zero-copy frame sharing between producer (streaming
gateway) and multiple consumers (inference servers).

Architecture:
- Producer creates SHM and writes frames in a ring pattern
- Consumers attach to SHM and read frames by index
- Overwrite is allowed (producer never waits for consumers)
- Consumers detect overwritten frames via frame_idx validation

Layout:
┌────────────────────────┐
│ Header (64 bytes)      │  write_idx, width, height, format, slot_count, ts_ns
├────────────────────────┤
│ Per-slot metadata      │  Per slot (16 bytes each):
│                        │    - frame_idx (8 bytes): monotonic frame index
│                        │    - seq_start (4 bytes): incremented BEFORE write
│                        │    - seq_end (4 bytes): incremented AFTER write
├────────────────────────┤
│ Frame slot 0           │  Fixed size: calculated from width, height, format
├────────────────────────┤
│ Frame slot 1           │
├────────────────────────┤
│ ...                    │
└────────────────────────┘

Torn Frame Detection:
- Producer increments seq_start BEFORE writing frame data
- Producer increments seq_end AFTER writing frame data
- Consumer reads seq_start, then frame data, then seq_end
- If seq_start != seq_end, the frame was being overwritten during read (torn)
"""

import logging
import struct
import time
from multiprocessing import shared_memory
from typing import Optional, Tuple, Union
import numpy as np


class ShmRingBuffer:
    """Shared memory ring buffer for raw frame storage.

    Supports NV12, RGB, and BGR frame formats for efficient video streaming.
    Uses a lock-free design where the producer overwrites old frames without
    waiting for consumers.

    Example (Producer):
        buffer = ShmRingBuffer(
            camera_id="cam_001",
            width=1920,
            height=1080,
            frame_format=ShmRingBuffer.FORMAT_BGR,  # Default - no conversion needed
            slot_count=30,
            create=True
        )
        frame_idx, slot = buffer.write_frame(bgr_frame.tobytes())

    Example (Consumer):
        buffer = ShmRingBuffer(
            camera_id="cam_001",
            width=1920,
            height=1080,
            frame_format=ShmRingBuffer.FORMAT_BGR,  # Default
            slot_count=30,
            create=False  # Attach to existing
        )
        if buffer.is_frame_valid(frame_idx):
            frame_data = buffer.read_frame(frame_idx)
    """

    # Header layout (64 bytes total, C-compatible for future GPU/CUDA interop)
    # Using struct format: '<QIIIIQ' = Little-endian
    #   Q = uint64 (8 bytes) - write_idx
    #   I = uint32 (4 bytes) - width
    #   I = uint32 (4 bytes) - height
    #   I = uint32 (4 bytes) - format
    #   I = uint32 (4 bytes) - slot_count
    #   Q = uint64 (8 bytes) - last_ts_ns (heartbeat timestamp)
    #   Padding to 64 bytes for alignment
    HEADER_FORMAT = '<QIIIIQ'
    HEADER_SIZE = 64  # Fixed header size (includes padding)

    # Per-slot metadata layout (16 bytes per slot):
    #   Q = uint64 (8 bytes) - frame_idx
    #   I = uint32 (4 bytes) - seq_start (incremented BEFORE write)
    #   I = uint32 (4 bytes) - seq_end (incremented AFTER write)
    SLOT_METADATA_SIZE = 16  # bytes per slot

    # Page alignment for hardware efficiency
    # Aligning frame slots to page boundaries improves:
    # - Memory mapping performance
    # - Cache efficiency
    # - Potential future DMA/GPU transfer efficiency
    PAGE_SIZE = 4096  # 4KB pages (standard on x86/x64)

    # Frame format constants
    FORMAT_NV12 = 0   # width * height * 1.5 bytes (GPU-friendly)
    FORMAT_RGB = 1    # width * height * 3 bytes
    FORMAT_BGR = 2    # width * height * 3 bytes (OpenCV native)

    # Format names for logging
    FORMAT_NAMES = {
        FORMAT_NV12: "NV12",
        FORMAT_RGB: "RGB",
        FORMAT_BGR: "BGR",
    }

    def __init__(
        self,
        camera_id: str,
        width: int,
        height: int,
        frame_format: int = FORMAT_BGR,
        slot_count: int = 30,
        create: bool = True,
    ):
        """Initialize SHM ring buffer.

        Args:
            camera_id: Unique camera identifier (used in SHM name)
            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 in ring buffer (default: 30)
            create: True for producer (creates SHM), False for consumer (attaches)

        Raises:
            ValueError: If frame_format is invalid
            FileExistsError: If create=True but SHM already exists (producer conflict)
            FileNotFoundError: If create=False but SHM doesn't exist (producer not started)
        """
        self.logger = logging.getLogger(f"{__name__}.ShmRingBuffer")

        # Validate inputs
        if frame_format not in self.FORMAT_NAMES:
            raise ValueError(f"Invalid frame_format: {frame_format}. Must be one of {list(self.FORMAT_NAMES.keys())}")
        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

        # Calculate frame size based on format
        self.frame_size = self._calculate_frame_size(width, height, frame_format)

        # Calculate page-aligned slot size for hardware efficiency
        self._aligned_slot_size = self._calculate_aligned_slot_size(self.frame_size)

        # Per-slot metadata: frame_idx + seq_start + seq_end (16 bytes per slot)
        self._slot_metadata_size = self.SLOT_METADATA_SIZE * slot_count

        # Calculate total SHM size (using aligned slots)
        self._total_size = self._calculate_total_size()

        # Windows-compatible SHM name (no slashes, max 255 chars)
        self.shm_name = self._generate_shm_name(camera_id)

        # Initialize shared memory
        self._shm: Optional[shared_memory.SharedMemory] = None
        self._init_shared_memory(create)

        # Calculate alignment overhead for logging
        alignment_overhead = self._aligned_slot_size - self.frame_size
        self.logger.info(
            f"ShmRingBuffer {'created' if create else 'attached'}: "
            f"name={self.shm_name}, size={self._total_size:,} bytes, "
            f"{width}x{height} {self.FORMAT_NAMES[frame_format]}, "
            f"{slot_count} slots, frame_size={self.frame_size:,} bytes, "
            f"aligned_slot_size={self._aligned_slot_size:,} bytes (+{alignment_overhead:,} padding)"
        )

    def _generate_shm_name(self, camera_id: str) -> str:
        """Generate Windows-compatible SHM name.

        Args:
            camera_id: Camera identifier

        Returns:
            Sanitized SHM name (max 200 chars, no special chars)
        """
        # Replace problematic characters and limit length
        safe_id = camera_id.replace("/", "_").replace("\\", "_").replace(":", "_")
        safe_id = "".join(c for c in safe_id if c.isalnum() or c == "_")
        return f"shm_cam_{safe_id[:180]}"

    def _calculate_frame_size(self, width: int, height: int, frame_format: int) -> int:
        """Calculate frame size in bytes based on format.

        Args:
            width: Frame width
            height: Frame height
            frame_format: Frame format constant

        Returns:
            Size in bytes for one frame
        """
        if frame_format == self.FORMAT_NV12:
            # NV12: Y plane (full) + UV plane (half height, interleaved)
            # Total = width * height * 1.5
            return int(width * height * 1.5)
        elif frame_format in (self.FORMAT_RGB, self.FORMAT_BGR):
            # RGB/BGR: 3 bytes per pixel
            return width * height * 3
        else:
            raise ValueError(f"Unknown frame format: {frame_format}")

    def _calculate_aligned_slot_size(self, frame_size: int) -> int:
        """Calculate page-aligned slot size for hardware efficiency.

        Pads frame slot size to PAGE_SIZE (4KB) boundaries for:
        - Better memory mapping performance
        - Improved cache efficiency
        - Future DMA/GPU transfer compatibility

        Args:
            frame_size: Unaligned frame size in bytes

        Returns:
            Page-aligned slot size in bytes
        """
        # Round up to next page boundary
        return ((frame_size + self.PAGE_SIZE - 1) // self.PAGE_SIZE) * self.PAGE_SIZE

    def _calculate_total_size(self) -> int:
        """Calculate total SHM size using page-aligned slots.

        Returns:
            Total size in bytes for entire SHM segment
        """
        # Use aligned slot size for frame data area
        return self.HEADER_SIZE + self._slot_metadata_size + (self._aligned_slot_size * self.slot_count)

    def _init_shared_memory(self, create: bool) -> None:
        """Initialize or attach to shared memory.

        Args:
            create: True to create new SHM, False to attach to existing

        Raises:
            FileExistsError: If create=True but SHM exists
            FileNotFoundError: If create=False but SHM doesn't exist
        """
        if create:
            try:
                # Try to create new SHM
                self._shm = shared_memory.SharedMemory(
                    name=self.shm_name,
                    create=True,
                    size=self._total_size
                )
                # 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)
                )
                # Initialize slot metadata (all zeros = no valid frames yet)
                for slot in range(self.slot_count):
                    self._write_slot_frame_idx(slot, 0)
                    self._write_slot_seq_start(slot, 0)
                    self._write_slot_seq_end(slot, 0)

            except FileExistsError:
                # On Windows, SHM may exist from previous crashed run
                # Try to unlink and recreate
                self.logger.warning(f"SHM {self.shm_name} exists, attempting to reuse/recreate")
                try:
                    # Attach and unlink old one
                    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}")

                # Try creating again
                self._shm = shared_memory.SharedMemory(
                    name=self.shm_name,
                    create=True,
                    size=self._total_size
                )
                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)
                )
                for slot in range(self.slot_count):
                    self._write_slot_frame_idx(slot, 0)
                    self._write_slot_seq_start(slot, 0)
                    self._write_slot_seq_end(slot, 0)
        else:
            # Consumer: attach to existing SHM
            try:
                self._shm = shared_memory.SharedMemory(name=self.shm_name, create=False)

                # Validate header matches expected configuration
                header = self._read_header()
                if header['width'] != self.width or header['height'] != self.height:
                    self.logger.warning(
                        f"SHM dimension mismatch: expected {self.width}x{self.height}, "
                        f"got {header['width']}x{header['height']}"
                    )
                if header['format'] != self.frame_format:
                    self.logger.warning(
                        f"SHM format mismatch: expected {self.FORMAT_NAMES.get(self.frame_format)}, "
                        f"got {self.FORMAT_NAMES.get(header['format'], 'unknown')}"
                    )

            except FileNotFoundError:
                raise FileNotFoundError(
                    f"SHM segment '{self.shm_name}' not found. "
                    f"Ensure producer (streaming gateway) is running first."
                )

    def _write_header(
        self,
        write_idx: int,
        width: int,
        height: int,
        frame_format: int,
        slot_count: int,
        last_ts_ns: int
    ) -> None:
        """Write header to SHM.

        Args:
            write_idx: Current write index
            width: Frame width
            height: Frame height
            frame_format: Frame format constant
            slot_count: Number of slots
            last_ts_ns: Timestamp in nanoseconds
        """
        header_bytes = struct.pack(
            self.HEADER_FORMAT,
            write_idx,
            width,
            height,
            frame_format,
            slot_count,
            last_ts_ns
        )
        # Pad to HEADER_SIZE
        header_bytes = header_bytes.ljust(self.HEADER_SIZE, b'\x00')
        self._shm.buf[:self.HEADER_SIZE] = header_bytes

    def _read_header(self) -> dict:
        """Read header from SHM.

        Returns:
            Dict with header fields: write_idx, width, height, format, slot_count, last_ts_ns
        """
        header_bytes = bytes(self._shm.buf[:struct.calcsize(self.HEADER_FORMAT)])
        write_idx, width, height, fmt, slot_count, last_ts_ns = struct.unpack(
            self.HEADER_FORMAT, header_bytes
        )
        return {
            'write_idx': write_idx,
            'width': width,
            'height': height,
            'format': fmt,
            'slot_count': slot_count,
            'last_ts_ns': last_ts_ns,
        }

    def _update_write_idx(self, new_idx: int) -> None:
        """Update write_idx in header (atomic uint64 write)."""
        # Write only the first 8 bytes (write_idx)
        self._shm.buf[:8] = struct.pack('<Q', new_idx)

    def _update_last_ts_ns(self, ts_ns: int) -> None:
        """Update heartbeat timestamp in header."""
        # last_ts_ns is at offset 24 (after write_idx=8 + width=4 + height=4 + format=4 + slot_count=4)
        offset = 8 + 4 + 4 + 4 + 4
        self._shm.buf[offset:offset+8] = struct.pack('<Q', ts_ns)

    def _get_slot_metadata_offset(self, slot: int) -> int:
        """Get offset for slot's metadata (frame_idx + seq_start + seq_end).

        Layout per slot (16 bytes):
            - frame_idx (8 bytes) at offset 0
            - seq_start (4 bytes) at offset 8
            - seq_end (4 bytes) at offset 12
        """
        return self.HEADER_SIZE + (slot * self.SLOT_METADATA_SIZE)

    def _write_slot_frame_idx(self, slot: int, frame_idx: int) -> None:
        """Write frame_idx to slot metadata."""
        offset = self._get_slot_metadata_offset(slot)
        self._shm.buf[offset:offset+8] = struct.pack('<Q', frame_idx)

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

    def _write_slot_seq_start(self, slot: int, seq: int) -> None:
        """Write seq_start counter for slot (incremented BEFORE frame write)."""
        offset = self._get_slot_metadata_offset(slot) + 8  # After frame_idx
        self._shm.buf[offset:offset+4] = struct.pack('<I', seq & 0xFFFFFFFF)

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

    def _write_slot_seq_end(self, slot: int, seq: int) -> None:
        """Write seq_end counter for slot (incremented AFTER frame write)."""
        offset = self._get_slot_metadata_offset(slot) + 12  # After frame_idx + seq_start
        self._shm.buf[offset:offset+4] = struct.pack('<I', seq & 0xFFFFFFFF)

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

    def _increment_slot_seq_start(self, slot: int) -> int:
        """Increment and return new seq_start for slot."""
        current = self._read_slot_seq_start(slot)
        new_seq = (current + 1) & 0xFFFFFFFF  # Wrap at 32-bit
        self._write_slot_seq_start(slot, new_seq)
        return new_seq

    def _increment_slot_seq_end(self, slot: int) -> int:
        """Increment and return new seq_end for slot."""
        current = self._read_slot_seq_end(slot)
        new_seq = (current + 1) & 0xFFFFFFFF  # Wrap at 32-bit
        self._write_slot_seq_end(slot, new_seq)
        return new_seq

    def _get_frame_offset(self, slot: int) -> int:
        """Get byte offset for frame data in given slot.

        Uses page-aligned slot size for hardware efficiency.
        """
        return self.HEADER_SIZE + self._slot_metadata_size + (slot * self._aligned_slot_size)

    def write_frame(self, raw_bytes: Union[bytes, memoryview, np.ndarray]) -> Tuple[int, int]:
        """Write frame to next slot (producer only).

        This method is NOT thread-safe - only one producer should write.
        Overwrites old frames in ring buffer pattern.

        Uses sequence counters for torn frame detection:
        - seq_start incremented BEFORE writing frame data
        - seq_end incremented AFTER writing frame data
        - Consumer can detect torn frames by checking seq_start != seq_end

        Args:
            raw_bytes: Raw frame data (NV12, RGB, or BGR bytes)

        Returns:
            Tuple of (frame_idx, slot_idx):
                - frame_idx: Monotonically increasing frame index
                - slot_idx: Physical slot where frame was written

        Raises:
            RuntimeError: If called on consumer instance
            ValueError: If raw_bytes size doesn't match expected frame_size
        """
        if not self._is_producer:
            raise RuntimeError("write_frame() can only be called on producer instance")

        # Handle numpy array input
        if isinstance(raw_bytes, np.ndarray):
            raw_bytes = raw_bytes.tobytes()
        elif isinstance(raw_bytes, memoryview):
            raw_bytes = bytes(raw_bytes)

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

        # Read current write_idx and increment
        header = self._read_header()
        frame_idx = header['write_idx'] + 1
        slot = frame_idx % self.slot_count

        # TORN FRAME DETECTION: Increment seq_start BEFORE writing frame data
        # This signals to consumers that a write is in progress
        self._increment_slot_seq_start(slot)

        # Write frame data to slot
        frame_offset = self._get_frame_offset(slot)
        self._shm.buf[frame_offset:frame_offset + self.frame_size] = raw_bytes

        # Update slot metadata (frame_idx for this slot)
        self._write_slot_frame_idx(slot, frame_idx)

        # TORN FRAME DETECTION: Increment seq_end AFTER writing frame data
        # When seq_start == seq_end, the frame write is complete
        self._increment_slot_seq_end(slot)

        # Update header: write_idx and timestamp (atomic-ish)
        ts_ns = int(time.time() * 1e9)
        self._update_write_idx(frame_idx)
        self._update_last_ts_ns(ts_ns)

        return frame_idx, slot

    def read_frame(self, frame_idx: int) -> Optional[memoryview]:
        """Read frame by index (consumer).

        Returns a memoryview into the shared memory for zero-copy access.
        Caller should copy the data if needed beyond the current frame.

        IMPORTANT: This returns a memoryview which may be overwritten by the producer.
        For safe reads, use read_frame_copy() instead.

        Args:
            frame_idx: Frame index to read

        Returns:
            memoryview of frame data, or None if frame was overwritten or torn
        """
        if not self.is_frame_valid(frame_idx):
            return None

        slot = frame_idx % self.slot_count
        frame_offset = self._get_frame_offset(slot)

        # Return memoryview for zero-copy access
        # Note: For memoryview, caller should verify frame wasn't torn after reading
        return self._shm.buf[frame_offset:frame_offset + self.frame_size]

    def read_frame_copy(self, frame_idx: int) -> Optional[bytes]:
        """Read frame and return a copy (consumer) with torn frame detection.

        Use this when you need the frame data to persist after
        the producer may overwrite the slot.

        This method detects torn frames by checking sequence counters:
        - Reads seq_start BEFORE reading frame data
        - Reads seq_end AFTER reading frame data
        - If seq_start != seq_end, the frame was overwritten during read (torn)

        Args:
            frame_idx: Frame index to read

        Returns:
            Bytes copy of frame data, or None if frame was overwritten or torn
        """
        if not self.is_frame_valid(frame_idx):
            return None

        slot = frame_idx % self.slot_count

        # TORN FRAME DETECTION: Read seq_start BEFORE reading frame data
        seq_start = self._read_slot_seq_start(slot)

        # Read frame data
        frame_offset = self._get_frame_offset(slot)
        frame_data = bytes(self._shm.buf[frame_offset:frame_offset + self.frame_size])

        # TORN FRAME DETECTION: Read seq_end AFTER reading frame data
        seq_end = self._read_slot_seq_end(slot)

        # If seq_start != seq_end, producer was writing during our read (torn frame)
        if seq_start != seq_end:
            self.logger.debug(
                f"Torn frame detected for frame_idx={frame_idx}, slot={slot}: "
                f"seq_start={seq_start}, seq_end={seq_end}"
            )
            return None

        # Also verify the frame_idx didn't change during read (double-check)
        stored_frame_idx = self._read_slot_frame_idx(slot)
        if stored_frame_idx != frame_idx:
            self.logger.debug(
                f"Frame overwritten during read: expected {frame_idx}, got {stored_frame_idx}"
            )
            return None

        return frame_data

    def is_frame_torn(self, frame_idx: int) -> bool:
        """Check if a frame read would be torn (producer writing during read).

        Args:
            frame_idx: Frame index to check

        Returns:
            True if the frame is currently being written (torn risk)
        """
        slot = frame_idx % self.slot_count
        seq_start = self._read_slot_seq_start(slot)
        seq_end = self._read_slot_seq_end(slot)
        return seq_start != seq_end

    def is_frame_valid(self, frame_idx: int) -> bool:
        """Check if frame_idx is still available (not overwritten).

        A frame is valid if:
        1. The slot's stored frame_idx matches the requested frame_idx
        2. The frame_idx is within the ring buffer window

        Args:
            frame_idx: Frame index to validate

        Returns:
            True if frame is still valid and readable
        """
        if frame_idx <= 0:
            return False

        # Check if within ring buffer window
        header = self._read_header()
        current_write_idx = header['write_idx']

        # Frame is too old if it's more than slot_count behind
        if current_write_idx - frame_idx >= self.slot_count:
            return False

        # Frame is in the future (hasn't been written yet)
        if frame_idx > current_write_idx:
            return False

        # Check slot metadata matches
        slot = frame_idx % self.slot_count
        stored_frame_idx = self._read_slot_frame_idx(slot)

        return stored_frame_idx == frame_idx

    def get_current_frame_idx(self) -> int:
        """Get latest written frame index.

        Returns:
            Current write_idx (0 if no frames written yet)
        """
        header = self._read_header()
        return header['write_idx']

    def get_last_heartbeat_ns(self) -> int:
        """Get last heartbeat timestamp in nanoseconds.

        Useful for detecting if producer is still alive.

        Returns:
            Last write timestamp in nanoseconds
        """
        header = self._read_header()
        return header['last_ts_ns']

    def get_header(self) -> dict:
        """Get full header information.

        Returns:
            Dict with all header fields
        """
        return self._read_header()

    def close(self) -> None:
        """Close and optionally unlink SHM.

        Producer unlinks (deletes) the SHM segment.
        Consumer just detaches without deleting.
        """
        if self._shm is None:
            return

        try:
            self._shm.close()

            if self._is_producer:
                # Producer is responsible for cleanup
                try:
                    self._shm.unlink()
                    self.logger.info(f"SHM segment {self.shm_name} unlinked")
                except FileNotFoundError:
                    pass  # Already unlinked (maybe by another process)
                except Exception as e:
                    self.logger.warning(f"Error unlinking SHM {self.shm_name}: {e}")
            else:
                self.logger.debug(f"Detached from SHM segment {self.shm_name}")

        except Exception as e:
            self.logger.warning(f"Error closing SHM {self.shm_name}: {e}")
        finally:
            self._shm = None

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()

    def __repr__(self) -> str:
        """String representation."""
        return (
            f"ShmRingBuffer(name={self.shm_name}, "
            f"{self.width}x{self.height} {self.FORMAT_NAMES.get(self.frame_format, 'unknown')}, "
            f"slots={self.slot_count}, producer={self._is_producer})"
        )


# =============================================================================
# NV12 Conversion Helper Functions
# =============================================================================

def bgr_to_nv12(bgr_frame: np.ndarray) -> bytes:
    """Convert BGR frame to NV12 format (GPU-friendly).

    NV12 layout:
    - Y plane: width * height bytes (luma - brightness)
    - UV plane: width * height / 2 bytes (interleaved chroma - color)
    - Total: width * height * 1.5 bytes

    This format is optimal for GPU-based inference (CUDA, TensorRT) as it:
    1. Matches camera sensor output (YUV native)
    2. Enables hardware-accelerated color space conversion
    3. Reduces memory bandwidth (1.5 vs 3 bytes per pixel)

    Args:
        bgr_frame: OpenCV BGR frame (numpy array, shape HxWx3, dtype uint8)

    Returns:
        NV12 bytes (Y plane followed by interleaved UV plane)

    Example:
        frame = cv2.imread("image.jpg")  # BGR format
        nv12_bytes = bgr_to_nv12(frame)
        assert len(nv12_bytes) == frame.shape[0] * frame.shape[1] * 1.5
    """
    import cv2

    height, width = bgr_frame.shape[:2]

    # BGR -> YUV_I420 (Y plane, U plane, V plane - all separate)
    yuv_i420 = cv2.cvtColor(bgr_frame, cv2.COLOR_BGR2YUV_I420)

    # I420 layout in OpenCV:
    # - Y plane: rows 0 to height-1 (full resolution)
    # - U plane: rows height to height + height/4 (quarter resolution)
    # - V plane: rows height + height/4 to height + height/2 (quarter resolution)

    # Extract planes
    y_end = height
    u_end = height + height // 4
    v_end = height + height // 2

    y_plane = yuv_i420[:y_end, :].flatten()
    u_plane = yuv_i420[y_end:u_end, :].flatten()
    v_plane = yuv_i420[u_end:v_end, :].flatten()

    # NV12 has interleaved UV (UVUVUV...) instead of separate U and V planes
    uv_interleaved = np.empty(len(u_plane) * 2, dtype=np.uint8)
    uv_interleaved[0::2] = u_plane
    uv_interleaved[1::2] = v_plane

    # Concatenate Y and interleaved UV
    return np.concatenate([y_plane, uv_interleaved]).tobytes()


def nv12_to_bgr(nv12_bytes: bytes, width: int, height: int) -> np.ndarray:
    """Convert NV12 bytes back to BGR frame.

    Args:
        nv12_bytes: NV12 format bytes
        width: Frame width
        height: Frame height

    Returns:
        BGR numpy array (shape HxWx3, dtype uint8)
    """
    import cv2

    # Calculate expected sizes
    y_size = width * height
    uv_size = y_size // 2
    expected_size = y_size + uv_size

    if len(nv12_bytes) != expected_size:
        raise ValueError(f"NV12 size mismatch: expected {expected_size}, got {len(nv12_bytes)}")

    # Create NV12 array for OpenCV
    # OpenCV expects NV12 in a specific layout
    nv12_array = np.frombuffer(nv12_bytes, dtype=np.uint8).reshape((height + height // 2, width))

    # Convert NV12 -> BGR
    bgr_frame = cv2.cvtColor(nv12_array, cv2.COLOR_YUV2BGR_NV12)

    return bgr_frame


def rgb_to_nv12(rgb_frame: np.ndarray) -> bytes:
    """Convert RGB frame to NV12 format.

    Args:
        rgb_frame: RGB frame (numpy array, shape HxWx3, dtype uint8)

    Returns:
        NV12 bytes
    """
    import cv2

    # RGB -> BGR -> NV12
    bgr_frame = cv2.cvtColor(rgb_frame, cv2.COLOR_RGB2BGR)
    return bgr_to_nv12(bgr_frame)
