"""Memory-Mapped Index Storage for RBSP.

Provides a persistent, memory-mapped key-value store backed by a single file.
Data survives process restarts and loads instantly via OS page-cache mapping.

File layout::

    [Header: 32 bytes]
        bytes 0-3:   Magic b"RIDX"
        bytes 4-7:   Version (uint32 LE)
        bytes 8-15:  Entry count (uint64 LE)
        bytes 16-23: Data region offset (uint64 LE) — start of KV data
        bytes 24-31: File size / next-write offset (uint64 LE)

    [Hash table: BUCKET_COUNT * 8 bytes]
        Each bucket is a uint64 LE offset into the data region (0 = empty).

    [Data region: variable]
        Each entry:
            key_len:   uint32 LE
            val_len:   uint32 LE
            key_bytes: key_len bytes (UTF-8 encoded)
            val_bytes: val_len bytes (raw)
            next_ptr:  uint64 LE (offset of next entry in chain, 0 = end)

All lookups are O(1) amortised via the hash table.
"""

from __future__ import annotations

import mmap
import os
import struct
import threading
from pathlib import Path
from typing import Iterator, Optional, Tuple

# --- Constants ---------------------------------------------------------------

MAGIC = b"RIDX"
VERSION = 1
HEADER_SIZE = 32
BUCKET_COUNT = 65536  # 64K buckets
BUCKET_TABLE_SIZE = BUCKET_COUNT * 8
DATA_START = HEADER_SIZE + BUCKET_TABLE_SIZE
INITIAL_FILE_SIZE = DATA_START + (1024 * 1024)  # 1 MiB initial data region
GROWTH_FACTOR = 2


# --- Header helpers ----------------------------------------------------------

def _pack_header(entry_count: int, data_offset: int, file_size: int) -> bytes:
    return struct.pack("<4sIQQQ", MAGIC, VERSION, entry_count, data_offset, file_size)


def _unpack_header(buf: bytes) -> Tuple[bytes, int, int, int, int]:
    return struct.unpack("<4sIQQQ", buf[:HEADER_SIZE])


# --- Entry helpers -----------------------------------------------------------

ENTRY_HEADER_FMT = "<II"  # key_len, val_len
ENTRY_HEADER_SIZE = struct.calcsize(ENTRY_HEADER_FMT)
NEXT_PTR_SIZE = 8  # uint64 LE


def _entry_total_size(key_len: int, val_len: int) -> int:
    return ENTRY_HEADER_SIZE + key_len + val_len + NEXT_PTR_SIZE


# --- Public API --------------------------------------------------------------

class MMapStore:
    """Persistent memory-mapped key-value store.

    Keys are UTF-8 strings.  Values are raw bytes (typically serialised
    ``NaturalSignature`` wire-format blobs).

    Thread-safe for concurrent reads; writes are serialised via a lock.

    Args:
        path:    Path to the backing file.  Created if it does not exist.
        readonly: Open in read-only mode (no writes allowed).
    """

    def __init__(self, path: str | Path, *, readonly: bool = False) -> None:
        self._path = Path(path)
        self._readonly = readonly
        self._lock = threading.Lock()
        self._mm: Optional[mmap.mmap] = None
        self._fd: Optional[int] = None
        self._file = None
        self._open()

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def _open(self) -> None:
        if self._path.exists():
            self._open_existing()
        elif self._readonly:
            raise FileNotFoundError(f"MMapStore file not found: {self._path}")
        else:
            self._create_new()

    def _create_new(self) -> None:
        """Create a fresh store file with empty header and bucket table."""
        self._path.parent.mkdir(parents=True, exist_ok=True)
        with open(self._path, "wb") as f:
            header = _pack_header(
                entry_count=0,
                data_offset=DATA_START,
                file_size=INITIAL_FILE_SIZE,
            )
            f.write(header)
            # Write empty bucket table
            f.write(b"\x00" * BUCKET_TABLE_SIZE)
            # Pad data region to initial size
            remaining = INITIAL_FILE_SIZE - DATA_START
            f.write(b"\x00" * remaining)
        self._open_existing()

    def _open_existing(self) -> None:
        """Memory-map an existing store file."""
        if self._readonly:
            self._file = open(self._path, "rb")
            self._mm = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_READ)
        else:
            self._file = open(self._path, "r+b")
            self._mm = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_WRITE)
        # Validate header
        magic, ver, count, data_off, fsize = _unpack_header(self._mm[:HEADER_SIZE])
        if magic != MAGIC:
            raise ValueError(f"Invalid MMapStore file: bad magic {magic!r}")
        if ver != VERSION:
            raise ValueError(f"Unsupported MMapStore version: {ver}")

    def close(self) -> None:
        """Flush and close the memory-mapped file."""
        if self._mm is not None:
            try:
                if not self._readonly:
                    self._mm.flush()
            except Exception:
                pass
            self._mm.close()
            self._mm = None
        if self._file is not None:
            self._file.close()
            self._file = None

    def __del__(self) -> None:
        try:
            self.close()
        except Exception:
            pass

    def __enter__(self) -> "MMapStore":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    # ------------------------------------------------------------------
    # Header access
    # ------------------------------------------------------------------

    def _read_entry_count(self) -> int:
        return struct.unpack_from("<Q", self._mm, 8)[0]

    def _write_entry_count(self, count: int) -> None:
        struct.pack_into("<Q", self._mm, 8, count)

    def _read_next_write_offset(self) -> int:
        """Read the file_size / next-write offset from the header."""
        return struct.unpack_from("<Q", self._mm, 24)[0]

    def _read_file_capacity(self) -> int:
        return struct.unpack_from("<Q", self._mm, 24)[0]

    def _write_next_write_offset(self, offset: int) -> None:
        struct.pack_into("<Q", self._mm, 24, offset)

    # ------------------------------------------------------------------
    # Bucket table
    # ------------------------------------------------------------------

    def _bucket_index(self, key: str) -> int:
        """Hash a key to a bucket index."""
        h = hash(key)
        return h % BUCKET_COUNT

    def _read_bucket(self, idx: int) -> int:
        """Read the head offset from bucket *idx*."""
        off = HEADER_SIZE + idx * 8
        return struct.unpack_from("<Q", self._mm, off)[0]

    def _write_bucket(self, idx: int, offset: int) -> None:
        off = HEADER_SIZE + idx * 8
        struct.pack_into("<Q", self._mm, off, offset)

    # ------------------------------------------------------------------
    # Entry read/write
    # ------------------------------------------------------------------

    def _read_entry_at(self, offset: int) -> Tuple[str, bytes, int]:
        """Read an entry at *offset*.  Returns ``(key, value, next_ptr)``."""
        key_len, val_len = struct.unpack_from(ENTRY_HEADER_FMT, self._mm, offset)
        kv_start = offset + ENTRY_HEADER_SIZE
        key = self._mm[kv_start : kv_start + key_len].decode("utf-8")
        val = bytes(self._mm[kv_start + key_len : kv_start + key_len + val_len])
        next_off = kv_start + key_len + val_len
        next_ptr = struct.unpack_from("<Q", self._mm, next_off)[0]
        return key, val, next_ptr

    def _grow_file(self, needed: int) -> None:
        """Grow the backing file to accommodate *needed* additional bytes."""
        current_cap = self._read_file_capacity()
        new_cap = max(current_cap * GROWTH_FACTOR, current_cap + needed + 4096)
        # Close mmap, resize file, re-open mmap
        self._mm.flush()
        self._mm.close()
        self._file.seek(0, 2)
        current_actual = self._file.tell()
        if new_cap > current_actual:
            self._file.write(b"\x00" * (new_cap - current_actual))
            self._file.flush()
        self._mm = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_WRITE)
        self._write_next_write_offset(new_cap)

    def _write_entry_at(
        self,
        offset: int,
        key: str,
        value: bytes,
        next_ptr: int = 0,
    ) -> int:
        """Write an entry at *offset*.  Returns the offset after this entry."""
        key_bytes = key.encode("utf-8")
        key_len = len(key_bytes)
        val_len = len(value)
        struct.pack_into(ENTRY_HEADER_FMT, self._mm, offset, key_len, val_len)
        kv_start = offset + ENTRY_HEADER_SIZE
        self._mm[kv_start : kv_start + key_len] = key_bytes
        self._mm[kv_start + key_len : kv_start + key_len + val_len] = value
        next_off = kv_start + key_len + val_len
        struct.pack_into("<Q", self._mm, next_off, next_ptr)
        return next_off + NEXT_PTR_SIZE

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def get(self, key: str) -> Optional[bytes]:
        """O(1) amortised lookup.  Returns the value bytes or ``None``."""
        with self._lock:
            bucket_idx = self._bucket_index(key)
            offset = self._read_bucket(bucket_idx)
            while offset != 0:
                k, v, next_ptr = self._read_entry_at(offset)
                if k == key:
                    return v
                offset = next_ptr
            return None

    def put(self, key: str, value: bytes) -> None:
        """Insert or update a key-value pair.

        Args:
            key:   UTF-8 string key.
            value: Raw bytes value.

        Raises:
            RuntimeError: If the store is opened read-only.
        """
        if self._readonly:
            raise RuntimeError("Store is read-only")

        with self._lock:
            bucket_idx = self._bucket_index(key)

            # Check if key already exists — overwrite by writing a new entry
            # and unlinking the old one (lazy: we leave dead entries).
            head_offset = self._read_bucket(bucket_idx)
            prev_offset = 0
            cur = head_offset
            while cur != 0:
                k, v, next_ptr = self._read_entry_at(cur)
                if k == key:
                    # Key exists — unlink from chain
                    if prev_offset == 0:
                        self._write_bucket(bucket_idx, next_ptr)
                    else:
                        # Patch previous entry's next_ptr
                        pk, pv, _ = self._read_entry_at(prev_offset)
                        self._write_entry_at(prev_offset, pk, pv, next_ptr)
                    count = self._read_entry_count()
                    self._write_entry_count(count - 1)
                    head_offset = self._read_bucket(bucket_idx)
                    break
                prev_offset = cur
                cur = next_ptr

            # Append new entry at end of data region
            write_off = self._get_write_cursor()
            entry_size = _entry_total_size(len(key.encode("utf-8")), len(value))
            capacity = self._read_file_capacity()
            if write_off + entry_size > capacity:
                self._grow_file(entry_size)

            # Write entry with next_ptr pointing to current bucket head
            new_head = self._read_bucket(bucket_idx)
            self._write_entry_at(write_off, key, value, new_head)
            self._write_bucket(bucket_idx, write_off)
            self._set_write_cursor(write_off + entry_size)

            count = self._read_entry_count()
            self._write_entry_count(count + 1)

    def delete(self, key: str) -> bool:
        """Remove a key.  Returns True if the key existed.

        Raises:
            RuntimeError: If the store is opened read-only.
        """
        if self._readonly:
            raise RuntimeError("Store is read-only")

        with self._lock:
            bucket_idx = self._bucket_index(key)
            prev_offset = 0
            cur = self._read_bucket(bucket_idx)
            while cur != 0:
                k, v, next_ptr = self._read_entry_at(cur)
                if k == key:
                    if prev_offset == 0:
                        self._write_bucket(bucket_idx, next_ptr)
                    else:
                        pk, pv, _ = self._read_entry_at(prev_offset)
                        self._write_entry_at(prev_offset, pk, pv, next_ptr)
                    count = self._read_entry_count()
                    self._write_entry_count(count - 1)
                    return True
                prev_offset = cur
                cur = next_ptr
            return False

    def __contains__(self, key: str) -> bool:
        return self.get(key) is not None

    def __len__(self) -> int:
        with self._lock:
            return self._read_entry_count()

    def keys(self) -> Iterator[str]:
        """Iterate over all keys (scans the bucket table)."""
        with self._lock:
            for bi in range(BUCKET_COUNT):
                offset = self._read_bucket(bi)
                while offset != 0:
                    k, v, next_ptr = self._read_entry_at(offset)
                    yield k
                    offset = next_ptr

    def items(self) -> Iterator[Tuple[str, bytes]]:
        """Iterate over all (key, value) pairs."""
        with self._lock:
            for bi in range(BUCKET_COUNT):
                offset = self._read_bucket(bi)
                while offset != 0:
                    k, v, next_ptr = self._read_entry_at(offset)
                    yield k, v
                    offset = next_ptr

    def flush(self) -> None:
        """Flush pending writes to disk."""
        if self._mm and not self._readonly:
            self._mm.flush()

    # ------------------------------------------------------------------
    # Write cursor — stored in a reserved region of the header
    # ------------------------------------------------------------------
    # We reuse the data_offset field (bytes 16-23) as the initial cursor,
    # and track the current cursor as the file_size field (bytes 24-31)
    # minus DATA_START tells us what's been used.  For simplicity, we
    # store the cursor as an in-memory value seeded from the file.

    _write_cursor: Optional[int] = None

    def _get_write_cursor(self) -> int:
        if self._write_cursor is None:
            # Scan to find end of data
            self._write_cursor = self._find_data_end()
        return self._write_cursor

    def _set_write_cursor(self, off: int) -> None:
        self._write_cursor = off

    def _find_data_end(self) -> int:
        """Walk all chains to find the maximum used offset."""
        max_end = DATA_START
        for bi in range(BUCKET_COUNT):
            offset = self._read_bucket(bi)
            while offset != 0:
                key_len, val_len = struct.unpack_from(ENTRY_HEADER_FMT, self._mm, offset)
                entry_end = offset + _entry_total_size(key_len, val_len)
                if entry_end > max_end:
                    max_end = entry_end
                kv_start = offset + ENTRY_HEADER_SIZE
                next_off = kv_start + key_len + val_len
                offset = struct.unpack_from("<Q", self._mm, next_off)[0]
        return max_end
