"""mgf.hrb._shm — the client side of the Part 5.3 shared-memory data plane.

The broker's ``open_shm_channel`` op allocates a private, memfd-backed
region, hosts a lock-free SPSC ring in it, and passes the region's fd via
``SCM_RIGHTS``. :class:`ShmChannel` maps that fd, validates the in-region
header against the broker's wire descriptor (never trusting the descriptor
blindly), and exposes the **producer** half of the ring — the client
produces; the broker drains via its asymmetric-trust consumer.

The byte layout is the ABI fixed in ``mgf-hrb-rt/src/ring.zig``
(``ShmRingHeader`` + ``HostedRing``). Little-endian; both peers are LE
Linux. Single-producer/single-consumer: the client writes ``head`` + the
slots, the broker writes ``tail``.
"""

from __future__ import annotations

import mmap
import os
import struct
from typing import TYPE_CHECKING, Any

from mgf.hrb.exceptions import MgfHrbError

if TYPE_CHECKING:
    from types import TracebackType

# ── ABI (must match mgf-hrb-rt/src/ring.zig) ──────────────────────────────
_MAGIC = 0x484D4252
_VERSION = 1
# Header at offset 0: magic u32, version u32, payload_size u32,
# slot_count u32, slots_offset u32, _reserved u32, capacity_mask u64.
_HEADER_FMT = "<IIIIIIQ"
# Slot: seq u64, payload u64 (the slice-1 payload type is u64).
_SLOT_FMT = "<QQ"
_SLOT_SIZE = struct.calcsize(_SLOT_FMT)  # 16
_U64_MASK = 0xFFFF_FFFF_FFFF_FFFF
_HEADER_SIZE = struct.calcsize(_HEADER_FMT)  # 32

# F-HRB-003 (asymmetric shm trust, Tier-0 #5) — the broker passes the region
# fd + a descriptor over the trusted socket, but the descriptor's sizes/offsets
# are still untrusted *numbers*: a buggy or compromised broker could name a
# region_bytes that blows up the mapping, or offsets that fall outside it.
# Bound them here so a bad descriptor yields a clean typed MgfHrbError instead
# of a raw ValueError from mmap or a struct.error from an out-of-range access.
# 256 MiB is far above any real ring (a few KiB) yet a firm DoS ceiling.
_MIN_REGION_BYTES = _HEADER_SIZE
_MAX_REGION_BYTES = 256 * 1024 * 1024


class ShmChannel:
    """A mapped shared-memory ring channel. Produce with :meth:`push`; the
    broker consumes via ``drain_shm_channel``. Close releases the mapping
    (the broker reclaims its side on disconnect)."""

    __slots__ = (
        "_closed",
        "_dropped_off",
        "_head_off",
        "_mask",
        "_mm",
        "_slot_count",
        "_slots_off",
        "_tail_off",
        "channel",
    )

    def __init__(self, fd: int, descriptor: dict[str, Any]) -> None:
        region_bytes = int(descriptor["region_bytes"])
        # F-HRB-003 — bound region_bytes BEFORE mmap. We own `fd`, so the
        # reject path must close it to avoid a leak.
        if not (_MIN_REGION_BYTES <= region_bytes <= _MAX_REGION_BYTES):
            os.close(fd)
            raise MgfHrbError(
                f"shm channel: region_bytes {region_bytes} outside the allowed "
                f"range [{_MIN_REGION_BYTES}, {_MAX_REGION_BYTES}]"
            )
        # Map the passed fd; the mapping keeps the pages alive, so the fd
        # itself can be closed immediately afterward.
        self._mm = mmap.mmap(
            fd,
            region_bytes,
            flags=mmap.MAP_SHARED,
            prot=mmap.PROT_READ | mmap.PROT_WRITE,
        )
        os.close(fd)
        self._closed = False
        try:
            self._validate(descriptor)
        except Exception:
            self._mm.close()
            self._closed = True
            raise
        self.channel = int(descriptor["channel"])
        self._slot_count = int(descriptor["slot_count"])
        self._mask = self._slot_count - 1
        self._slots_off = int(descriptor["slots_offset"])
        self._head_off = int(descriptor["head_offset"])
        self._tail_off = int(descriptor["tail_offset"])
        self._dropped_off = int(descriptor["dropped_offset"])

    def _validate(self, descriptor: dict[str, Any]) -> None:
        """The descriptor is the broker's word; the header is ground truth.
        They must agree, and the layout must be one we understand."""
        magic, version, payload_size, slot_count, slots_offset, _res, cap_mask = struct.unpack_from(
            _HEADER_FMT, self._mm, 0
        )
        if magic != _MAGIC:
            raise MgfHrbError(f"shm channel: bad magic 0x{magic:08x} (expected 0x{_MAGIC:08x})")
        if version != _VERSION:
            raise MgfHrbError(f"shm channel: unsupported version {version}")
        if payload_size != 8:
            raise MgfHrbError(f"shm channel: unsupported payload_size {payload_size}")
        if (
            payload_size != int(descriptor["payload_size"])
            or slot_count != int(descriptor["slot_count"])
            or slots_offset != int(descriptor["slots_offset"])
            or cap_mask != slot_count - 1
        ):
            raise MgfHrbError("shm channel: in-region header disagrees with descriptor")

        # F-HRB-003 — slot_count must be a positive power of two: `idx & mask`
        # (push()) only wraps correctly when mask == slot_count - 1 is a full
        # low-bit mask. The header-agreement check above ties cap_mask to
        # slot_count - 1 but does not by itself force power-of-two-ness.
        if slot_count <= 0 or (slot_count & (slot_count - 1)) != 0:
            raise MgfHrbError(
                f"shm channel: slot_count {slot_count} is not a positive power of two"
            )

        # F-HRB-003 — every field this client reads/writes must lie fully
        # inside the mapped region. Without this, a descriptor naming an
        # offset past region_bytes turns push()'s struct access into a raw
        # struct.error (or, on a larger mapping, a silent read of unrelated
        # pages). Bound each span end against region_bytes; reject negatives.
        region_bytes = int(descriptor["region_bytes"])
        offsets = {
            "slots_offset": (slots_offset, slot_count * _SLOT_SIZE),
            "head_offset": (int(descriptor["head_offset"]), 8),
            "tail_offset": (int(descriptor["tail_offset"]), 8),
            "dropped_offset": (int(descriptor["dropped_offset"]), 8),
        }
        for name, (off, span) in offsets.items():
            if off < 0 or off + span > region_bytes:
                raise MgfHrbError(
                    f"shm channel: {name}={off} (+{span}) falls outside the "
                    f"{region_bytes}-byte region"
                )

    @property
    def slot_count(self) -> int:
        return self._slot_count

    def push(self, item: int) -> bool:
        """Publish one u64 into the ring (producer half). Returns False if
        the ring is full (drop_newest — the broker sees a seq gap). Writes
        the payload+seq, then bumps ``head`` last so the broker never reads
        a slot before its payload is in place."""
        if self._closed:
            raise MgfHrbError("shm channel is closed")
        (head,) = struct.unpack_from("<Q", self._mm, self._head_off)
        (tail,) = struct.unpack_from("<Q", self._mm, self._tail_off)
        if (head - tail) & _U64_MASK >= self._slot_count:
            # Ring full: drop_newest. Bump the shared dropped counter so the
            # broker's drain reports honest accounting (a seq gap, Part 5.3.3).
            (dropped,) = struct.unpack_from("<Q", self._mm, self._dropped_off)
            struct.pack_into("<Q", self._mm, self._dropped_off, (dropped + 1) & _U64_MASK)
            return False
        idx = head & self._mask
        slot_off = self._slots_off + idx * _SLOT_SIZE
        struct.pack_into(_SLOT_FMT, self._mm, slot_off, head, item)  # seq=head, payload
        struct.pack_into("<Q", self._mm, self._head_off, (head + 1) & _U64_MASK)
        return True

    def close(self) -> None:
        if not self._closed:
            self._closed = True
            self._mm.close()

    def __enter__(self) -> ShmChannel:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        self.close()
