"""Reference solution — Write-Ahead Log and Crash Recovery.

DO NOT READ BEFORE YOU HAVE RUN THE PROBLEM UNDER THE CLOCK.
Teaching text: ../../../WARMUP.md § Chapter 8.

write() returning tells you the KERNEL accepted the bytes. It says nothing
about the disk. A process crash is survivable at that point; a machine crash is
not. Only fsync pushes toward stable media -- and even then you are trusting
the device not to lie about its volatile cache.

Framing: [4-byte length][payload][4-byte CRC32]. The CRC is the part people
skip and it is the important one. Without it, a torn write whose length field
happened to be complete reads back as a VALID-LOOKING record full of garbage.
That is silent corruption, which is strictly worse than a crash. The checksum
converts it into a clean truncation.

Checkpointing must itself be crash-safe: write to a temp file, fsync the DATA,
os.replace (atomic on POSIX), then fsync the DIRECTORY -- because the rename is
a directory-entry change and that entry also has to reach stable storage. That
last fsync is the one almost everybody forgets.
"""

from __future__ import annotations

import os
import struct
import zlib

HEADER = struct.Struct("<I")
CRC = struct.Struct("<I")
CHECKPOINT_HEADER = struct.Struct("<QI")     # log offset, blob length


class WriteAheadLog:
    """Append-only log with CRC framing, torn-tail recovery, and checkpoints."""

    def __init__(self, path, fsync_policy="always", group_size=64,
                 fsync=os.fsync):
        if fsync_policy not in ("always", "group", "never"):
            raise ValueError("fsync_policy must be always, group or never")
        self.path = path
        self._policy = fsync_policy
        self._group_size = group_size
        self._fsync = fsync
        self._pending = 0
        self.fsync_count = 0
        self.valid_end = 0
        self._file = open(path, "a+b")
        self._file.seek(0, os.SEEK_END)
        self.offset = self._file.tell()

    # ---- appending -------------------------------------------------------
    def append(self, payload: bytes) -> int:
        record = HEADER.pack(len(payload)) + payload + CRC.pack(zlib.crc32(payload))
        self._file.write(record)
        self.offset += len(record)
        self._pending += 1
        if self._policy == "always":
            self._durable()
        elif self._policy == "group" and self._pending >= self._group_size:
            # Group commit: one disk round trip covers many records. This is
            # how real databases escape one fsync per commit.
            self._durable()
        return self.offset

    def _durable(self):
        self._file.flush()
        if self._policy != "never":
            self._fsync(self._file.fileno())
            self.fsync_count += 1
        self._pending = 0

    def flush(self):
        self._file.flush()
        if self._policy != "never" and self._pending:
            self._fsync(self._file.fileno())
            self.fsync_count += 1
        self._pending = 0

    def close(self):
        self.flush()
        self._file.close()

    # ---- replay ----------------------------------------------------------
    def replay(self, from_offset=0):
        """Yield (end_offset, payload). A bad tail ENDS the iteration."""
        with open(self.path, "rb") as handle:
            handle.seek(from_offset)
            position = from_offset
            while True:
                head = handle.read(HEADER.size)
                if len(head) < HEADER.size:
                    break                              # (1) no length prefix
                (length,) = HEADER.unpack(head)
                body = handle.read(length)
                if len(body) < length:
                    break                              # (2) incomplete payload
                tail = handle.read(CRC.size)
                if len(tail) < CRC.size:
                    break
                (expected,) = CRC.unpack(tail)
                if zlib.crc32(body) != expected:
                    break                              # (3) torn or corrupt
                position += HEADER.size + length + CRC.size
                yield position, body
            self.valid_end = position

    def truncate_to_valid(self) -> int:
        for _ in self.replay():
            pass
        self._file.flush()
        self._file.close()
        with open(self.path, "r+b") as handle:
            handle.truncate(self.valid_end)
        self._file = open(self.path, "a+b")
        self._file.seek(0, os.SEEK_END)
        self.offset = self._file.tell()
        return self.valid_end

    # ---- checkpointing ---------------------------------------------------
    @property
    def _checkpoint_path(self):
        return self.path + ".ckpt"

    def checkpoint(self, blob: bytes) -> int:
        """Durable and atomic: either the old checkpoint or the new one."""
        self.flush()
        payload = CHECKPOINT_HEADER.pack(self.offset, len(blob)) + blob
        payload += CRC.pack(zlib.crc32(payload))

        tmp = self._checkpoint_path + ".tmp"
        with open(tmp, "wb") as handle:
            handle.write(payload)
            handle.flush()
            self._fsync(handle.fileno())      # 1. the DATA is durable
            self.fsync_count += 1
        os.replace(tmp, self._checkpoint_path)  # 2. atomic on POSIX

        directory = os.path.dirname(os.path.abspath(self._checkpoint_path))
        dir_fd = os.open(directory, os.O_DIRECTORY)
        try:
            self._fsync(dir_fd)               # 3. the DIRECTORY ENTRY is durable
            self.fsync_count += 1
        finally:
            os.close(dir_fd)
        return self.offset

    def load_checkpoint(self):
        """Return (blob, log_offset), or None if absent or corrupt."""
        path = self._checkpoint_path
        if not os.path.exists(path):
            return None
        with open(path, "rb") as handle:
            raw = handle.read()
        if len(raw) < CHECKPOINT_HEADER.size + CRC.size:
            return None
        body, tail = raw[:-CRC.size], raw[-CRC.size:]
        (expected,) = CRC.unpack(tail)
        if zlib.crc32(body) != expected:
            return None                        # a half-written checkpoint
        log_offset, length = CHECKPOINT_HEADER.unpack(
            body[:CHECKPOINT_HEADER.size])
        blob = body[CHECKPOINT_HEADER.size:CHECKPOINT_HEADER.size + length]
        if len(blob) != length:
            return None
        return blob, log_offset