"""
Shared memory implementation of a ringbuffer.

This is designed for multiple writers and a single reader.  A semaphore is used to
lock the data strocture so that only one process can read or write to it at once.
If the ring buffer is full, the writer will overwrite the oldest messages to make room
for the new message.

This is designed so that the ring buffer lives outside of any particular process,
meaning either the reader, writer or both can be stopped or restarted without losing
queued messages.  This means the shared memory is not cleaned up by anyone; you can
manually clean it up like this:
 rm /dev/shm/{name}.
However that only works properly in Python 3.13+.  In lower Python versions, exiting
a process will delete the shared memory.

Semaphore Timeout Protection:
By default, semaphore acquisition uses a 5-second timeout to protect against orphaned
locks (when a process crashes while holding the semaphore). If a timeout occurs, a
TimeoutError is raised with instructions for manual recovery. You can disable timeouts
by passing lock_timeout=None to the constructor.
"""
import posix_ipc
import asyncio
import struct
from concurrent.futures import ThreadPoolExecutor
import time

try:
    from multiprocessing import shared_memory
    SHARED_MEMORY_AVAILABLE = True
except ImportError:
    # Python 3.7 doesn't have shared_memory in multiprocessing
    shared_memory = None
    SHARED_MEMORY_AVAILABLE = False


class RingBuffer:
    META_SIZE = 12  # Metadata: head (4 bytes), tail (4 bytes), capacity (4 bytes)
    EXISTS_CACHE_TIME = 60  # cache exists result

    def __init__(self, name, capacity=None, create=False, lock_timeout=5.0):
        self.name = name
        self.capacity = capacity
        self.create = create
        self.lock_timeout = lock_timeout  # Timeout in seconds for lock acquisition (None = block forever)
        self.shm = None
        self.lock = None
        self.executor = ThreadPoolExecutor() if SHARED_MEMORY_AVAILABLE else None
        self._exists = None  # cached result
        self._check_time = 0

    def exists(self):
        """Check if the ringbuffer shared memory exists."""
        if not SHARED_MEMORY_AVAILABLE:
            return False
        now = time.monotonic()
        if self._exists is not None and now < self._check_time:
            return self._exists
        try:
            try:
                shm = shared_memory.SharedMemory(name=self.name, track=False)
            except TypeError:
                shm = shared_memory.SharedMemory(name=self.name)
            shm.close()
            self._exists = True
        except FileNotFoundError:
            self._exists = False
            self.close()
        self._check_time = now + self.EXISTS_CACHE_TIME
        return self._exists

    def _create(self):
        """Create shared memory and semaphore."""
        size = self.META_SIZE + self.capacity  # Total shared memory size
        try:
            self.shm = shared_memory.SharedMemory(name=self.name,
                                                  create=self.create,
                                                  size=size,
                                                  track=False)
        except TypeError:  # can't control tracking in python < 3.13
            self.shm = shared_memory.SharedMemory(name=self.name,
                                                  create=self.create,
                                                  size=size)
        struct.pack_into('iii', self.shm.buf, 0, 0, 0,
                         self.capacity)  # Initialize head, tail, capacity
        self.lock = posix_ipc.Semaphore(self.name, posix_ipc.O_CREAT, initial_value=1)

    def _attach(self):
        """Create or attach shared memory"""
        if self.create:
            try:
                self._create()
                return
            except FileExistsError:
                pass  # already created, attach instead

        # Attach instead of creating
        try:
            self.shm = shared_memory.SharedMemory(name=self.name, track=False)
        except TypeError:
            self.shm = shared_memory.SharedMemory(name=self.name)
        except FileNotFoundError:
            self._exists = False
            self.close()
            return
        self.lock = posix_ipc.Semaphore(self.name)

        # Initialize capacity
        if not self.capacity:
            self._acquire_lock()
            try:
                _, _, self.capacity = self._get_metadata()
            finally:
                self.lock.release()

    def _acquire_lock(self):
        """Acquire the lock with timeout. Raises TimeoutError if lock cannot be acquired."""
        if self.lock_timeout is None:
            # Block forever (original behavior)
            self.lock.acquire()
        else:
            try:
                self.lock.acquire(self.lock_timeout)
            except posix_ipc.BusyError:
                raise TimeoutError(
                    f"Could not acquire lock '{self.name}' within {self.lock_timeout}s. "
                    "The semaphore may be orphaned (previous process crashed while holding it). "
                    "Consider manually cleaning up: rm /dev/shm/sem.{self.name}")

    def close(self):
        """Close the shared memory and shutdown executor."""
        if self.shm:
            self.shm.close()
        if self.executor:
            self.executor.shutdown(wait=False)
        self.shm = self.lock = None
        self.executor = None

    def unlink(self):
        """Unlink shared memory and semaphore."""
        try:
            if self.shm:
                self.shm.unlink()
        except FileNotFoundError:
            pass  # Already unlinked
        try:
            if self.lock:
                self.lock.unlink()
        except posix_ipc.ExistentialError:
            pass  # Already unlinked

    def _get_metadata(self):
        """Retrieve head, tail, and capacity (assumes lock is held by caller)."""
        return struct.unpack_from('iii', self.shm.buf, 0)

    def _set_metadata(self, head, tail):
        """Update head and tail pointers (assumes lock is held by caller)."""
        struct.pack_into('ii', self.shm.buf, 0, head, tail)

    def _two_phase_write(self, start: int, buf: bytes) -> None:
        """Write `buf` into the ring buffer in up to two phases if it crosses the boundary."""
        ring_start = self.META_SIZE
        ring_end = self.META_SIZE + self.capacity  # One past last valid byte
        length = len(buf)

        # Phase 1: write up to the end of the buffer
        space_to_end = ring_end - start
        if length <= space_to_end:
            # Fits in one contiguous slice
            self.shm.buf[start:start + length] = buf
        else:
            # Partially fill the end
            chunk1_size = space_to_end
            self.shm.buf[start:ring_end] = buf[:chunk1_size]

            # Phase 2: wrap to the front
            chunk2_size = length - chunk1_size
            wrap_start = ring_start  # typically the start offset of ring data
            self.shm.buf[wrap_start:wrap_start + chunk2_size] = buf[chunk1_size:]

    def _two_phase_read(self, start: int, length: int) -> bytes:
        """Read `length` bytes from the ring buffer in up to two phases if it crosses the boundary."""
        ring_start = self.META_SIZE
        ring_end = self.META_SIZE + self.capacity
        data_out = bytearray(length)

        space_to_end = ring_end - start
        if length <= space_to_end:
            # Single chunk read
            data_out[:] = self.shm.buf[start:start + length]
        else:
            # Phase 1: read until the end
            chunk1_size = space_to_end
            data_out[:chunk1_size] = self.shm.buf[start:ring_end]

            # Phase 2: wrap to the front
            chunk2_size = length - chunk1_size
            wrap_start = ring_start
            data_out[chunk1_size:] = self.shm.buf[wrap_start:wrap_start + chunk2_size]

        return bytes(data_out)

    def _write(self, data: bytes):
        """Blocking: Write a serialized message to the buffer, discarding old messages if needed."""
        if self.lock is None:
            self._attach()
        self._acquire_lock()
        try:
            head, tail, capacity = self._get_metadata()
            message_size = len(data)
            total_size = message_size + 4  # 4 bytes for the length prefix

            if total_size > capacity:
                raise ValueError("Message too large for buffer")

            # Keep discarding old messages until there's enough space
            next_tail = (tail + total_size) % capacity
            while (next_tail - head) % capacity < total_size and head != tail:
                # Read old message size at head
                prefix_start = self.META_SIZE + head
                old_size_bytes = self._two_phase_read(prefix_start, 4)
                old_message_size = struct.unpack('>I', old_size_bytes)[0]

                head = (head + old_message_size + 4) % capacity
                # Recompute next_tail after discarding
                _, tail_, _ = self._get_metadata()
                tail = tail_
                next_tail = (tail + total_size) % capacity

            # Now write the combined prefix + data in two-phase
            full_msg = struct.pack('>I', message_size) + data
            start_offset = self.META_SIZE + tail
            self._two_phase_write(start_offset, full_msg)

            # Update metadata
            self._set_metadata(head, next_tail)
        finally:
            self.lock.release()

    def _read(self):
        """Blocking: Read a serialized message from the buffer, using two-phase read if needed."""
        if self.lock is None:
            self._attach()
        self._acquire_lock()
        try:
            head, tail, capacity = self._get_metadata()
            if head == tail:
                return None  # Buffer empty

            # First read the 4-byte prefix
            prefix_start = self.META_SIZE + head
            prefix_bytes = self._two_phase_read(prefix_start, 4)
            message_size = struct.unpack('>I', prefix_bytes)[0]

            # Validate message size to prevent reading corrupted data
            if message_size == 0 or message_size > capacity:
                raise ValueError(
                    f"Corrupted ring buffer: invalid message_size={message_size}, capacity={capacity}"
                )

            # Then read the message itself
            message_start = (head + 4) % capacity
            data = self._two_phase_read(self.META_SIZE + message_start, message_size)

            # Advance head
            next_head = (head + message_size + 4) % capacity
            self._set_metadata(next_head, tail)

            return data
        finally:
            self.lock.release()

    async def get_metadata(self):
        """Async wrapper for retrieving metadata."""
        loop = asyncio.get_running_loop()

        def _get():
            self._acquire_lock()
            try:
                return self._get_metadata()
            finally:
                self.lock.release()

        return await loop.run_in_executor(self.executor, _get)

    async def write(self, data: bytes):
        """Async wrapper for writing data."""
        loop = asyncio.get_running_loop()
        try:
            await loop.run_in_executor(self.executor, self._write, data)
        except FileNotFoundError:
            self._exists = False
            self.close()

    async def read(self):
        """Async wrapper for reading data."""
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(self.executor, self._read)
