"""POSIX shared-memory ring writer compatible with the C++ IPC ring."""

from __future__ import annotations

import logging
import mmap
import os
from pathlib import Path
from typing import Final

from python_app.models.run_config_schema import RingEndpointModel
from python_app.models.run_config_validation import validate_ring_endpoint
from python_app.orchestration.shm import atomics
logger = logging.getLogger(__name__)

_HEADER_SIZE: Final[int] = 64
_SLOT_HEADER_SIZE: Final[int] = 16
_MAGIC: Final[bytes] = b"RDRRING2"
_VERSION: Final[int] = 1


class ShmRingWriter:
    """Write binary payloads into the shared-memory ring used by C++ workers.

    The writer is the sole *owner* of the rings it opens: there is exactly one
    producer per ring (the acquisition producer for the raw/raw_tap rings). On a
    geometry mismatch with a pre-existing segment (e.g. a stale ring left by a prior
    run with a different sweep config), the owner unlinks and recreates the segment
    from scratch rather than truncating in place or diverging silently — mirroring
    the clean-shm-on-restart contract on the C++/deploy side (#13). A non-owner must
    never recreate a ring; readers and C++ consumers only ever attach to an existing
    one.
    """

    def __init__(self, ring_name: str, capacity: int, slot_size_bytes: int) -> None:
        """Open or create a POSIX SHM ring by name (as the ring owner)."""
        if not ring_name.startswith("/") or len(ring_name) < 2 or "/" in ring_name[1:] or "\x00" in ring_name:
            raise ValueError("ring_name must be '/' followed by a single nonempty filename")
        validate_ring_endpoint(RingEndpointModel(ring_name, capacity, slot_size_bytes))

        self._ring_name = ring_name
        self._capacity = int(capacity)
        self._slot_size_bytes = int(slot_size_bytes)
        self._mapped_size = _HEADER_SIZE + self._capacity * (_SLOT_HEADER_SIZE + self._slot_size_bytes)
        self._path = Path("/dev/shm") / ring_name[1:]
        self._file = self._mmap = None
        try:
            self._open_owned()
        except BaseException:
            self.close()
            raise

    def _open_owned(self) -> None:
        """Open the ring, recreating it from scratch on a geometry/header mismatch.

        As the single owner of this ring we may safely discard a stale segment: a
        size or header mismatch means the existing segment belongs to an earlier,
        incompatible run, so we unlink it and create a fresh one instead of mapping
        an inconsistent layout.
        """
        created = not self._path.exists()
        fd = os.open(self._path, os.O_RDWR | os.O_CREAT, 0o660)
        self._file = os.fdopen(fd, "r+b", buffering=0)
        if created or self._path.stat().st_size != self._mapped_size:
            # Wrong-sized stale segment: drop it entirely and recreate, so the file
            # and any future mapping agree on geometry instead of being truncated
            # under a producer/consumer that still expects the old layout.
            if not created:
                logger.info("Ring %s: stale segment with wrong size, recreating", self._ring_name)
                self._file.close()
                self._unlink_if_present()
                fd = os.open(self._path, os.O_RDWR | os.O_CREAT | os.O_EXCL, 0o660)
                self._file = os.fdopen(fd, "r+b", buffering=0)
            self._file.truncate(self._mapped_size)
            created = True

        self._mmap = mmap.mmap(self._file.fileno(), self._mapped_size)
        if created:
            self._initialize_header()
            logger.debug(
                "Created SHM ring writer %s (capacity=%d, slot_size=%d bytes)",
                self._ring_name,
                self._capacity,
                self._slot_size_bytes,
            )
            return

        # Size matched but the header geometry/magic does not: the owner recreates
        # rather than diverge. Unlink and reopen as a brand-new ring.
        if not self._header_matches():
            logger.warning("Ring %s: header/geometry mismatch on existing segment, recreating", self._ring_name)
            self._mmap.close()
            self._file.close()
            self._unlink_if_present()
            fd = os.open(self._path, os.O_RDWR | os.O_CREAT | os.O_EXCL, 0o660)
            self._file = os.fdopen(fd, "r+b", buffering=0)
            self._file.truncate(self._mapped_size)
            self._mmap = mmap.mmap(self._file.fileno(), self._mapped_size)
            self._initialize_header()
        else:
            logger.debug("Reusing existing SHM ring writer %s", self._ring_name)

    def _unlink_if_present(self) -> None:
        """Remove the backing /dev/shm file if it exists (owner-only operation)."""
        try:
            self._path.unlink()
        except FileNotFoundError:
            pass

    def close(self) -> None:
        """Close mmap and file handle."""
        if self._mmap is not None:
            self._mmap.close()
            self._mmap = None
        if self._file is not None:
            self._file.close()
            self._file = None
        logger.debug("Closed SHM ring writer %s", self._ring_name)

    def push(self, payload: bytes) -> bool:
        """Push one payload with overwrite-oldest semantics on overflow."""
        if len(payload) > self._slot_size_bytes:
            return False

        write_seq = self._read_u64(24)
        read_seq = self._read_u64(32)
        if max(0, write_seq - read_seq) >= self._capacity:
            # Ring full: the consumer is not keeping up, so the oldest unread slot is
            # overwritten (overwrite-oldest). Logged at DEBUG to avoid flooding the
            # log when a backlog persists across many pushes.
            logger.debug("Ring %s: full, overwriting oldest unread slot", self._ring_name)
            # Advance the consumer cursor past the slot we are about to overwrite, but
            # re-read it first and move it only forward: a concurrent reader may have
            # already advanced it, and clobbering that backward would re-deliver an
            # already-consumed slot as a duplicate.
            atomics.advance(self._mmap, 32, read_seq + 1)
            self._write_u64(40, self._read_u64(40) + 1)

        index = write_seq % self._capacity
        slot_offset = _HEADER_SIZE + index * (_SLOT_HEADER_SIZE + self._slot_size_bytes)
        payload_offset = slot_offset + _SLOT_HEADER_SIZE

        # Invalidate the previous publication before changing its bytes. A peek or
        # a lapped consumer must not accept the old sequence during this overwrite.
        self._write_u64(slot_offset + 8, 0)
        atomics.fence()
        self._write_u32(slot_offset, len(payload))
        self._write_u32(slot_offset + 4, 0)
        self._mmap[payload_offset : payload_offset + len(payload)] = payload
        self._write_u64(slot_offset + 8, write_seq + 1)
        self._write_u64(24, write_seq + 1)
        return True

    @property
    def ring_name(self) -> str:
        """Return the POSIX SHM ring name."""
        return self._ring_name

    @property
    def slot_size_bytes(self) -> int:
        """Return maximum payload size per slot."""
        return self._slot_size_bytes

    def _initialize_header(self) -> None:
        # A new ftruncate-extended segment is already zero-filled by the kernel;
        # avoid allocating and copying another complete ring just to initialize it.
        self._mmap[:_HEADER_SIZE] = b"\x00" * _HEADER_SIZE
        self._mmap[:8] = _MAGIC
        self._write_u32(8, _VERSION)
        self._write_u32(12, self._capacity)
        self._write_u32(16, self._slot_size_bytes)
        self._write_u32(20, 0)
        self._write_u64(24, 0)
        self._write_u64(32, 0)
        self._write_u64(40, 0)

    def _header_matches(self) -> bool:
        """Return whether the existing segment's header matches this ring's geometry.

        Used by the owner to decide whether a same-sized pre-existing segment can be
        reused as-is or must be unlinked and recreated.
        """
        return (
            self._mmap[:8] == _MAGIC
            and self._read_u32(8) == _VERSION
            and self._read_u32(12) == self._capacity
            and self._read_u32(16) == self._slot_size_bytes
        )

    def _read_u32(self, offset: int) -> int:
        return atomics.load(self._mmap, offset, 4)

    def _read_u64(self, offset: int) -> int:
        return atomics.load(self._mmap, offset, 8)

    def _write_u32(self, offset: int, value: int) -> None:
        atomics.store(self._mmap, offset, int(value), 4)

    def _write_u64(self, offset: int, value: int) -> None:
        atomics.store(self._mmap, offset, int(value), 8)