"""A key-value store whose durability comes from a write-ahead log.

Every mutation is appended to a log file BEFORE the in-memory map is
touched. That ordering is the "write-ahead" in write-ahead log, and it is
what makes recovery possible: the log is the truth, the map is a cache of
it that happens to live in RAM. Recovery is a replay from byte zero.

Two knobs, `flush` and `fsync`, expose the two buffers a write passes
through on its way to the platter:

    kv.put()  ->  Python's buffer  ->  kernel page cache  ->  disk
                                 flush()             fsync()

They are the whole point of the toy. They decide which crashes you
survive, and one of them is completely invisible to `kill -9`.
"""

import os
import struct
import zlib

# Each record is framed as: payload length | crc32 of payload | payload.
# The length is what lets recovery find the next record; the checksum is
# what lets it notice that the last one is a lie.
_HEADER = struct.Struct(">II")

_PUT = b"P"
_DELETE = b"D"


def _encode(op, key, value=b""):
    """op byte | 2-byte key length | key | value (rest of the payload)."""
    return op + struct.pack(">H", len(key)) + key + value


def _decode(payload):
    op = payload[:1]
    (klen,) = struct.unpack_from(">H", payload, 1)
    return op, payload[3:3 + klen], payload[3 + klen:]


def _frame(payload):
    return _HEADER.pack(len(payload), zlib.crc32(payload)) + payload


class WalKV:
    """put/get/delete, backed by an append-only log.

    `flush=False` keeps bytes in Python's own buffer, where a SIGKILL
    destroys them. `fsync=True` forces them all the way to disk and is the
    only setting that survives losing power. `fsync` implies `flush` --
    os.fsync works on a file descriptor and knows nothing about the
    userspace buffer sitting above it.
    """

    def __init__(self, path, flush=True, fsync=False):
        self.path = path
        self.flush = flush or fsync
        self.fsync = fsync
        self.data = {}

        # Recovery statistics, so a caller can see what the log gave back.
        self.replayed = 0    # records successfully applied
        self.discarded = 0   # trailing bytes that failed framing or checksum

        self._recover()
        self._f = open(path, "ab")

    # -- writes ---------------------------------------------------------

    def put(self, key, value):
        self._append(_encode(_PUT, key, value))
        self.data[key] = value

    def delete(self, key):
        # A delete is a record too -- a tombstone. Without one, replaying
        # the log would faithfully resurrect every key you ever removed.
        self._append(_encode(_DELETE, key))
        self.data.pop(key, None)

    def get(self, key, default=None):
        return self.data.get(key, default)

    def _append(self, payload):
        # Log first, mutate second. The reverse order would let a crash
        # land between them and leave the map ahead of its own log.
        self._f.write(_frame(payload))
        if self.flush:
            self._f.flush()            # Python's buffer -> kernel page cache
        if self.fsync:
            os.fsync(self._f.fileno())  # kernel page cache -> disk

    # -- recovery -------------------------------------------------------

    def _recover(self):
        """Replay the log, stopping at the first record that isn't intact.

        A crash can tear the final record in half, so the last few bytes of
        a log are always suspect. Stopping is deliberate: everything before
        the damage is still perfectly good, and refusing to start because
        the tail is ragged would throw away a healthy database.
        """
        if not os.path.exists(self.path):
            return

        with open(self.path, "rb") as f:
            blob = f.read()

        off = 0
        while off + _HEADER.size <= len(blob):
            length, crc = _HEADER.unpack_from(blob, off)
            start = off + _HEADER.size
            payload = blob[start:start + length]

            if len(payload) < length:          # truncated mid-record
                break
            if zlib.crc32(payload) != crc:     # framing intact, bytes rotten
                break

            self._apply(payload)
            off = start + length
            self.replayed += 1

        self.discarded = len(blob) - off

    def _apply(self, payload):
        op, key, value = _decode(payload)
        if op == _PUT:
            self.data[key] = value
        else:
            self.data.pop(key, None)

    def close(self):
        self._f.close()


def simulate_power_loss(path, fsynced_bytes):
    """Model a power cut, which is the one failure you cannot stage.

    Killing a process leaves the kernel -- and therefore the page cache --
    completely untouched, so no signal can ever show you this. Pulling the
    plug empties that cache. What remains is exactly what fsync() had
    already pushed to the platter; everything after it never existed.
    """
    with open(path, "r+b") as f:
        f.truncate(fsynced_bytes)