from inspect import signature
from multiprocessing.shared_memory import SharedMemory
from struct import calcsize, pack, unpack
from typing import Optional, Self
from math import ceil

from posix_ipc import Semaphore, O_CREX

from edri.config.setting import CHUNK_SIZE

_SHM_SUPPORTS_TRACK = "track" in signature(SharedMemory).parameters


class SharedMemoryPipe:
    """
    A bounded, single-producer/single-consumer byte pipe implemented on top of
    POSIX shared memory + POSIX semaphores.

    The shared memory region is divided into `max_slots` fixed-size slots.
    Each slot contains a small header followed by up to `slot_size - HEADER_SIZE`
    bytes of payload.

    Synchronization:
      - `slots` semaphore counts free slots (writer acquires before writing).
      - `items` semaphore counts filled slots (reader acquires before reading).

    Termination:
      - EOF is signaled by writing a slot whose header has `eof=1`.
      - `read()` returns None on EOF; the iterator stops.

    Notes:
      - Intended for local IPC between processes on the same host.
      - The writer end owns creation of SHM and semaphores; the reader end
        attaches by name.
    """
    HEADER_FMT = "I B 3x"  # uint32 size, uint8 eof, 3-byte pad
    HEADER_SIZE = calcsize(HEADER_FMT)

    _ALIGNMENT = 64  # bytes – tune here if you like

    @staticmethod
    def _aligned(n: int, alignment: int = _ALIGNMENT) -> int:
        """Round `n` up to the next multiple of `alignment` bytes."""
        return ((n + alignment - 1) // alignment) * alignment

    @classmethod
    def _compute_slot_size(cls,
                           *,
                           total_size: int | None,
                           max_slots: int,
                           min_slot_size: int) -> int:
        """
        Choose an efficient slot size.

        If `total_size` is provided, this derives a slot size large enough that
        `total_size` bytes can be transported in at most `max_slots` slots.

        The computed size:
          - reserves space for the per-slot header,
          - is rounded up to `_ALIGNMENT` bytes (cache friendliness),
          - is never smaller than `min_slot_size`.

        If `total_size` is None, returns `min_slot_size` unchanged.
        """
        if total_size is None:
            return min_slot_size

        # Bytes of **payload** we must be able to carry per slot
        payload_per_slot = ceil(total_size / max_slots)

        # Add header, align, and honour the caller’s minimum
        candidate = cls._aligned(payload_per_slot + cls.HEADER_SIZE)
        return max(min_slot_size, candidate)

    def __init__(self,
                 name: Optional[str] = None,
                 *,
                 max_slots: int = 64,
                 slot_size: int = CHUNK_SIZE,
                 total_size: int | None = None,
                 _closed: bool = False,):
        """
        Create a shared-memory-backed pipe endpoint.

        This class has two modes:

        **Writer mode** (default):
          - `name` is None
          - creates a new shared memory segment sized `slot_size * max_slots`
          - creates two POSIX semaphores:
              - `<shm.name>_items` (initial 0): number of filled slots
              - `<shm.name>_slots` (initial `max_slots`): number of free slots

        **Reader mode**:
          - `name` is the shared memory name created by the writer
          - attaches to the existing shared memory segment and semaphores

        Parameters
        ----------
        name:
            Shared memory name to attach to. If None, a new shared memory region is
            created and this instance becomes the writer end.
        max_slots:
            Number of fixed-size slots in the ring buffer. Higher values increase
            buffering and reduce contention, at the cost of more shared memory.
        slot_size:
            Minimum slot size in bytes (including header). Each slot holds a
            `HEADER_SIZE` header plus up to `slot_size - HEADER_SIZE` payload bytes.
            If `total_size` is provided, this acts as a lower bound only.
        total_size:
            Optional expected total number of payload bytes to be sent. If provided,
            the constructor may increase `slot_size` (aligned to `_ALIGNMENT`) so
            that the full payload can be transmitted in at most `max_slots` chunks,
            reducing chunk count and semaphore traffic.
        _closed:
            Internal flag used for pickling / state restoration. When True, no OS
            resources are created or attached.

        Attributes set
        --------------
        name:
            The shared memory name (generated in writer mode, provided in reader mode).
        is_writer:
            True if this endpoint created the shared memory (writer mode).
        slot_size:
            Final slot size actually used (may be larger than the `slot_size` argument
            when `total_size` is provided).
        max_slots:
            Number of slots in the ring buffer.
        local_index:
            Slot cursor for this endpoint (write index for writer, read index for reader).

        Notes
        -----
        - On Python versions where `SharedMemory` supports the `track` parameter,
          this implementation disables resource tracking via `track=False`.
        - This implementation assumes a single writer and a single reader, each with
          its own `local_index`.
        """
        self.name = name
        self.max_slots = max_slots
        self.total_size = total_size
        self.slot_size = self._compute_slot_size(total_size=total_size,
                                                 max_slots=max_slots,
                                                 min_slot_size=slot_size)
        self.is_writer = self.name is None
        self._closed = _closed
        self.local_index = 0

        shm_kwargs = {}
        if _SHM_SUPPORTS_TRACK:
            shm_kwargs["track"] = False

        if self._closed:
            return
        if self.is_writer:
            shm_kwargs.update({"create": True, "size": self.slot_size * self.max_slots})
            self.shm = SharedMemory(**shm_kwargs)
            self.name = self.shm.name

            self._items_sem_name = self.shm.name + "_items"
            self._slots_sem_name = self.shm.name + "_slots"
            self.items = Semaphore(self._items_sem_name, flags=O_CREX, initial_value=0)
            self.slots = Semaphore(self._slots_sem_name, flags=O_CREX, initial_value=self.max_slots)
        else:
            shm_kwargs["name"] = self.name
            self._items_sem_name = self.name + "_items"
            self._slots_sem_name = self.name + "_slots"
            self.shm = SharedMemory(**shm_kwargs)
            self.items = Semaphore(self._items_sem_name)
            self.slots = Semaphore(self._slots_sem_name)

    def _write_slot(self, data: bytes, eof: bool):
        """
               Write a single slot worth of data into the ring buffer.

               This method:
                 1) waits for a free slot (`slots.acquire()`),
                 2) writes the slot header (payload length + EOF flag),
                 3) writes the payload bytes,
                 4) advances the local ring index,
                 5) signals availability to the reader (`items.release()`).

               Parameters
               ----------
               data:
                   Payload bytes to store in the slot. Must be <= max payload size
                   (`slot_size - HEADER_SIZE`).
               eof:
                   Whether this slot marks end-of-stream. When `eof` is True, the
                   reader will return None from `read()` after consuming this slot.

               Notes
               -----
               This is an internal primitive; `write()` handles chunking automatically.
               """
        self.slots.acquire()

        slot_index = self.local_index % self.max_slots
        offset = slot_index * self.slot_size

        header = pack(self.HEADER_FMT, len(data), int(eof))
        self.shm.buf[offset: offset + self.HEADER_SIZE] = header
        self.shm.buf[offset + self.HEADER_SIZE: offset + self.HEADER_SIZE + len(data)] = data

        self.local_index += 1
        self.items.release()

    def write(self, data: bytes, /, *, close: bool = False):
        """
        Write bytes into the pipe, splitting into slot-sized chunks as needed.

        Parameters
        ----------
        data:
            The bytes to send.
        close:
            If True, the final chunk of this call is marked as EOF, so the
            reader will observe end-of-stream after consuming it.

        Notes
        -----
        - Chunk size is `slot_size - HEADER_SIZE`.
        - If `close=True` and `data` is empty, no chunks are written; use
          `close()` to reliably send EOF in all cases.
        """
        max_payload = self.slot_size - self.HEADER_SIZE
        total_chunks = (len(data) + max_payload - 1) // max_payload

        for i in range(total_chunks):
            chunk = data[i * max_payload: (i + 1) * max_payload]
            eof = close and (i == total_chunks - 1)
            self._write_slot(chunk, eof=eof)

    def read(self) -> Optional[bytes]:
        """
        Read the next chunk from the pipe.

        Returns
        -------
        bytes | None
            The next payload chunk, or None if an EOF slot was encountered.

        Behavior
        --------
        This method:
          1) waits for an available item (`items.acquire()`),
          2) reads the header to determine payload length and EOF flag,
          3) copies out the payload bytes,
          4) advances the local ring index,
          5) frees the slot (`slots.release()`).

        Notes
        -----
        Returning None is the EOF sentinel used by `__next__` to stop iteration.
        """
        self.items.acquire()

        slot_index = self.local_index % self.max_slots
        offset = slot_index * self.slot_size

        header = self.shm.buf[offset: offset + self.HEADER_SIZE]
        size, eof = unpack(self.HEADER_FMT, header)

        data = bytes(self.shm.buf[offset + self.HEADER_SIZE
                                  : offset + self.HEADER_SIZE + size])
        self.local_index += 1

        self.slots.release()
        return None if eof else data

    def reader(self) -> Self:
        """
        Create a reader endpoint attached to the same underlying shared memory.

        Returns
        -------
        SharedMemoryPipe
            A new instance configured as the reader end of this pipe.

        Raises
        ------
        RuntimeError
            If called on an instance that is already a reader.

        Notes
        -----
        The returned reader shares the same SHM name and semaphore names,
        but maintains its own local read index.
        """
        if not self.is_writer:
            raise RuntimeError("Only writer can spawn reader")
        return SharedMemoryPipe(name=self.name,
                                max_slots=self.max_slots,
                                slot_size=self.slot_size,
                                total_size=self.total_size)

    def close(self):
        """
        Close this endpoint.

        Writer semantics:
          - Enqueues an explicit EOF slot (empty payload with eof=1),
            then closes the shared memory handle.

        Reader semantics:
          - Closes the shared memory handle and unlinks the shared memory
            segment (removing it from the system namespace).

        Notes
        -----
        - This method is idempotent.
        - Depending on your lifecycle expectations, you may prefer that the
          writer is responsible for unlinking SHM instead of the reader.
        """
        if self._closed:
            return
        self._closed = True

        if self.is_writer:
            # send EOF downstream
            self._write_slot(b"", eof=True)
            self.shm.close()
        else:
            self.shm.close()
            self.shm.unlink()

        self.shm = None  # guard against accidental reuse

    def __enter__(self):
        return self

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

    def __iter__(self) -> Self:
        """
        Return an iterator over incoming chunks (reader side only).

        Raises
        ------
        RuntimeError
            If called on the writer end.
        """
        if self.is_writer:
            raise RuntimeError("Cannot iterate over the writer end of the pipe")
        return self

    def __next__(self) -> bytes:
        """
        Yield the next chunk from the pipe.

        On EOF, this closes the reader endpoint and raises StopIteration.
        """
        if self._closed:
            raise StopIteration

        chunk = self.read()
        if chunk is None:
            self.close()
            raise StopIteration
        return chunk

    def __getstate__(self):
        """
        Support pickling of the reader endpoint.

        The writer cannot be pickled because it owns system resources and is
        responsible for emitting EOF and coordinating lifetime.

        Returns
        -------
        dict
            Minimal state needed to reattach to the shared memory region and
            semaphores in another process.

        Raises
        ------
        RuntimeError
            If invoked on the writer end.
        """
        if self.is_writer:
            raise RuntimeError("Writer cannot be pickled")
        return {
            "name": self.name,
            "max_slots": self.max_slots,
            "slot_size": self.slot_size,
            "total_size": self.total_size,
            "_closed": self._closed,
        }

    def __setstate__(self, state):
        """
        Restore a pickled reader endpoint by reattaching to SHM/semaphores.

        Parameters
        ----------
        state:
            The dictionary produced by `__getstate__`.
        """
        self.__init__(**state)
