"""Write-Ahead Log (WAL) for RBSP.

Records all index mutations (insert, delete, update, segment operations) before
they are applied to in-memory structures.  On crash recovery, the WAL is
replayed from the last checkpoint to restore consistent state.

Zero external dependencies -- uses only stdlib.

WAL file header (16 bytes)::

    [4 bytes: magic "RWAL"]
    [4 bytes: version uint32 big-endian]
    [8 bytes: starting sequence number uint64 big-endian]

Each WAL entry on disk::

    [4 bytes: sentinel 0xFEEDFACE uint32 big-endian]
    [4 bytes: total_length uint32 big-endian]
    [1 byte:  operation enum]
    [8 bytes: timestamp float64 big-endian]
    [4 bytes: data_length uint32 big-endian]
    [N bytes: data (variable, opaque)]
    [4 bytes: CRC32 uint32 big-endian over op+timestamp+data_len+data]

Checkpoint marker file::

    [8 bytes: sequence number uint64 big-endian]

Thread-safe: all public methods acquire an internal lock.
"""

from __future__ import annotations

import enum
import os
import struct
import threading
import time
import zlib
from dataclasses import dataclass
from pathlib import Path
from typing import IO, List, Optional, Tuple


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

WAL_MAGIC = b"RWAL"
WAL_VERSION = 1
WAL_HEADER_SIZE = 16  # magic(4) + version(4) + sequence_start(8)
WAL_HEADER_FMT = ">4sIQ"

_ENTRY_SENTINEL = 0xFEEDFACE
_ENTRY_HEADER_FMT = ">IIBdI"  # sentinel(4) + total_len(4) + op(1) + timestamp(8) + data_len(4)
_ENTRY_HEADER_SIZE = struct.calcsize(_ENTRY_HEADER_FMT)  # 21 bytes
_CRC_FMT = ">I"
_CRC_SIZE = 4

_CHECKPOINT_FILENAME = "checkpoint"


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class WALOperation(enum.IntEnum):
    """Operation types recorded in WAL entries."""

    INSERT = 1
    DELETE = 2
    UPDATE = 3
    CHECKPOINT = 4
    SEGMENT_CREATE = 5
    SEGMENT_MERGE = 6
    SEGMENT_DELETE = 7


class SyncMode(enum.Enum):
    """Durability modes for WAL writes.

    Attributes:
        FSYNC:     ``os.fsync`` after every append (safest, slowest).
        FDATASYNC: ``os.fdatasync`` where available, falls back to fsync.
        NONE:      No sync -- rely on OS buffer cache (fastest, least durable).
    """

    FSYNC = "fsync"
    FDATASYNC = "fdatasync"
    NONE = "none"


# ---------------------------------------------------------------------------
# Data Classes
# ---------------------------------------------------------------------------

@dataclass(frozen=True, slots=True)
class WALConfig:
    """Configuration for the write-ahead log.

    Attributes:
        log_dir:             Directory for WAL and checkpoint files.
        sync_mode:           Durability mode for writes.
        max_log_size:        Maximum WAL file size in bytes before rotation.
        checkpoint_interval: Number of entries between auto-checkpoints.
    """

    log_dir: str
    sync_mode: SyncMode = SyncMode.FSYNC
    max_log_size: int = 64 * 1024 * 1024  # 64 MiB
    checkpoint_interval: int = 1000


@dataclass(frozen=True, slots=True)
class WALEntry:
    """A single WAL entry.

    Attributes:
        sequence:  Monotonically increasing sequence number.
        operation: The WAL operation type.
        timestamp: POSIX timestamp when the entry was written.
        data:      Opaque payload bytes (typically JSON-encoded).
        checksum:  CRC32 checksum for integrity verification.
    """

    sequence: int
    operation: WALOperation
    timestamp: float
    data: bytes
    checksum: int


# ---------------------------------------------------------------------------
# Entry Packing / Unpacking
# ---------------------------------------------------------------------------

def _pack_entry(operation: WALOperation, timestamp: float, data: bytes) -> Tuple[bytes, int]:
    """Pack a WAL entry into its binary wire format.

    Args:
        operation: The operation type.
        timestamp: POSIX timestamp.
        data:      Opaque payload bytes.

    Returns:
        Tuple of (packed_bytes, crc32_checksum).
    """
    crc_payload = struct.pack(">BdI", int(operation), timestamp, len(data)) + data
    crc = zlib.crc32(crc_payload) & 0xFFFFFFFF
    total_length = 1 + 8 + 4 + len(data) + 4  # op + timestamp + data_len + data + crc
    packed = (
        struct.pack(">II", _ENTRY_SENTINEL, total_length)
        + crc_payload
        + struct.pack(_CRC_FMT, crc)
    )
    return packed, crc


def _unpack_entry(data: bytes, offset: int, sequence: int) -> Tuple[Optional[WALEntry], int]:
    """Unpack a single WAL entry starting at *offset*.

    Args:
        data:     Raw bytes of the WAL file.
        offset:   Byte offset to start reading from.
        sequence: Sequence number to assign to this entry.

    Returns:
        Tuple of (entry_or_None, next_offset).  Returns ``(None, offset)``
        if the data is truncated, the sentinel is wrong, or CRC check fails.
    """
    remaining = len(data) - offset

    # Need at least sentinel(4) + total_length(4) = 8 bytes
    if remaining < 8:
        return None, offset

    sentinel, total_length = struct.unpack_from(">II", data, offset)
    if sentinel != _ENTRY_SENTINEL:
        return None, offset

    # total_length covers: op(1) + timestamp(8) + data_len(4) + data(N) + crc(4)
    if remaining < 8 + total_length:
        return None, offset  # truncated entry

    content_start = offset + 8
    content = data[content_start : content_start + total_length]

    # Minimum content: op(1) + timestamp(8) + data_len(4) + crc(4) = 17 bytes
    if total_length < 17:
        return None, offset

    op_byte = content[0]
    timestamp = struct.unpack_from(">d", content, 1)[0]
    data_len = struct.unpack_from(">I", content, 9)[0]

    # Validate that data_len is consistent with total_length
    expected_total = 1 + 8 + 4 + data_len + 4
    if expected_total != total_length:
        return None, offset

    entry_data = content[13 : 13 + data_len]
    stored_crc = struct.unpack_from(_CRC_FMT, content, 13 + data_len)[0]

    # CRC covers: op + timestamp + data_len + data
    crc_payload = content[: 13 + data_len]
    computed_crc = zlib.crc32(crc_payload) & 0xFFFFFFFF
    if stored_crc != computed_crc:
        return None, offset  # CRC mismatch

    try:
        operation = WALOperation(op_byte)
    except ValueError:
        return None, offset  # unknown operation

    entry = WALEntry(
        sequence=sequence,
        operation=operation,
        timestamp=timestamp,
        data=entry_data,
        checksum=stored_crc,
    )
    return entry, offset + 8 + total_length


# ---------------------------------------------------------------------------
# WAL Header Helpers
# ---------------------------------------------------------------------------

def _pack_header(sequence_start: int) -> bytes:
    """Pack a WAL file header."""
    return struct.pack(WAL_HEADER_FMT, WAL_MAGIC, WAL_VERSION, sequence_start)


def _unpack_header(data: bytes) -> Optional[Tuple[int, int]]:
    """Unpack a WAL file header.

    Returns:
        ``(version, sequence_start)`` or ``None`` if invalid.
    """
    if len(data) < WAL_HEADER_SIZE:
        return None
    magic, version, seq_start = struct.unpack_from(WAL_HEADER_FMT, data, 0)
    if magic != WAL_MAGIC:
        return None
    return version, seq_start


# ---------------------------------------------------------------------------
# Write-Ahead Log
# ---------------------------------------------------------------------------

class WriteAheadLog:
    """Write-ahead log for crash-recoverable index mutations.

    All index-modifying operations are first written to the WAL before
    being applied to in-memory structures.  On recovery, entries since
    the last checkpoint are replayed to restore consistent state.

    Thread-safe: all public methods acquire an internal lock.

    Args:
        config: WAL configuration.

    Usage::

        wal = WriteAheadLog(WALConfig(log_dir="/tmp/rbsp-wal"))
        seq = wal.append(WALOperation.INSERT, data=b'{"doc_id":"abc"}')
        wal.checkpoint()
        wal.close()

    Recovery::

        wal = WriteAheadLog(WALConfig(log_dir="/tmp/rbsp-wal"))
        entries = wal.recover()
        for entry in entries:
            # replay each operation
            ...
    """

    def __init__(self, config: WALConfig) -> None:
        self._config = config
        self._lock = threading.Lock()
        self._sequence: int = 0
        self._entries_since_checkpoint: int = 0
        self._current_log: Optional[IO[bytes]] = None
        self._current_log_path: Optional[Path] = None
        self._current_log_size: int = 0
        self._closed: bool = False

        self._log_dir = Path(config.log_dir)
        self._log_dir.mkdir(parents=True, exist_ok=True)

        self._initialize()

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def current_sequence(self) -> int:
        """The next sequence number that will be assigned."""
        with self._lock:
            return self._sequence

    @property
    def entries_since_checkpoint(self) -> int:
        """Number of entries written since the last checkpoint."""
        with self._lock:
            return self._entries_since_checkpoint

    # ------------------------------------------------------------------
    # Initialization
    # ------------------------------------------------------------------

    def _initialize(self) -> None:
        """Find existing WAL files or create a new one."""
        wal_files = self._get_wal_files()
        if wal_files:
            # Open the latest WAL file for appending
            latest = wal_files[-1]
            self._current_log_path = latest
            data = latest.read_bytes()
            result = _unpack_header(data)
            if result is not None:
                _version, seq_start = result
                # Count entries in this file to determine current sequence
                self._sequence = seq_start
                offset = WAL_HEADER_SIZE
                while offset < len(data):
                    entry, new_offset = _unpack_entry(data, offset, self._sequence)
                    if entry is None:
                        break
                    self._sequence += 1
                    offset = new_offset
            # Open for appending
            self._current_log = open(latest, "ab")
            self._current_log_size = latest.stat().st_size
        else:
            self._new_log_file()

    def _get_wal_files(self) -> List[Path]:
        """Return all WAL files in the log directory, sorted by name."""
        files = sorted(self._log_dir.glob("wal_*.log"))
        return files

    def _new_log_file(self) -> None:
        """Create a new WAL log file with a fresh header."""
        filename = f"wal_{self._sequence:016d}.log"
        path = self._log_dir / filename
        self._current_log_path = path

        # Write header
        with open(path, "wb") as f:
            f.write(_pack_header(self._sequence))
            f.flush()

        # Reopen for appending
        self._current_log = open(path, "ab")
        self._current_log_size = WAL_HEADER_SIZE

    # ------------------------------------------------------------------
    # Sync
    # ------------------------------------------------------------------

    def _sync(self) -> None:
        """Sync the WAL file to disk based on the configured sync mode."""
        if self._current_log is None:
            return

        self._current_log.flush()

        if self._config.sync_mode == SyncMode.FSYNC:
            os.fsync(self._current_log.fileno())
        elif self._config.sync_mode == SyncMode.FDATASYNC:
            fd = self._current_log.fileno()
            if hasattr(os, "fdatasync"):
                os.fdatasync(fd)
            else:
                os.fsync(fd)  # fallback on Windows
        # SyncMode.NONE: flush() only, no sync call

    # ------------------------------------------------------------------
    # Write Operations
    # ------------------------------------------------------------------

    def append(self, operation: WALOperation, data: bytes = b"") -> int:
        """Append an entry to the WAL.

        Args:
            operation: The operation type to record.
            data:      Opaque payload bytes (e.g. JSON-encoded dict).

        Returns:
            The sequence number assigned to this entry.

        Raises:
            RuntimeError: If the WAL has been closed.
        """
        with self._lock:
            if self._closed:
                raise RuntimeError("WAL is closed")

            timestamp = time.time()
            packed, crc = _pack_entry(operation, timestamp, data)

            seq = self._sequence
            self._current_log.write(packed)
            self._sync()

            self._sequence += 1
            self._current_log_size += len(packed)
            self._entries_since_checkpoint += 1

            # Auto-rotate if max_log_size exceeded
            if self._current_log_size >= self._config.max_log_size:
                self._rotate_internal()

        return seq

    # ------------------------------------------------------------------
    # Checkpoint
    # ------------------------------------------------------------------

    def checkpoint(self, snapshot_data: bytes = b"") -> int:
        """Create a checkpoint at the current position.

        Writes a CHECKPOINT entry to the WAL and updates the checkpoint
        marker file with the current sequence number.

        Args:
            snapshot_data: Optional snapshot metadata to include in the
                           checkpoint entry payload.

        Returns:
            The sequence number of the checkpoint entry.

        Raises:
            RuntimeError: If the WAL has been closed.
        """
        with self._lock:
            if self._closed:
                raise RuntimeError("WAL is closed")

            timestamp = time.time()
            packed, crc = _pack_entry(WALOperation.CHECKPOINT, timestamp, snapshot_data)

            seq = self._sequence
            self._current_log.write(packed)
            self._sync()

            self._sequence += 1
            self._current_log_size += len(packed)

            # Atomically write checkpoint marker file
            marker_path = self._log_dir / _CHECKPOINT_FILENAME
            tmp_path = marker_path.with_suffix(".tmp")
            with open(tmp_path, "wb") as f:
                f.write(struct.pack(">Q", seq))
                f.flush()
                os.fsync(f.fileno())
            os.replace(str(tmp_path), str(marker_path))

            self._entries_since_checkpoint = 0

        return seq

    def _read_checkpoint_sequence(self) -> int:
        """Read the checkpoint sequence number from the marker file.

        Returns:
            The checkpoint sequence number, or 0 if no checkpoint exists.
        """
        marker_path = self._log_dir / _CHECKPOINT_FILENAME
        if not marker_path.exists():
            return 0
        try:
            data = marker_path.read_bytes()
            if len(data) < 8:
                return 0
            (seq,) = struct.unpack(">Q", data[:8])
            return seq
        except (OSError, struct.error):
            return 0

    # ------------------------------------------------------------------
    # Recovery
    # ------------------------------------------------------------------

    def recover(self) -> List[WALEntry]:
        """Replay WAL entries after the last checkpoint.

        Reads all WAL files in order and returns entries whose sequence
        number is greater than or equal to the checkpoint sequence.  Entries
        with bad CRC32 checksums or truncated data are silently skipped
        (they represent incomplete writes from a crash).

        Returns:
            Ordered list of WALEntry objects to replay.
        """
        with self._lock:
            checkpoint_seq = self._read_checkpoint_sequence()
            wal_files = self._get_wal_files()

            entries: List[WALEntry] = []
            max_seq = self._sequence

            for wal_file in wal_files:
                try:
                    file_data = wal_file.read_bytes()
                except OSError:
                    continue

                result = _unpack_header(file_data)
                if result is None:
                    continue  # skip files with invalid headers

                _version, seq_start = result
                current_seq = seq_start
                offset = WAL_HEADER_SIZE

                while offset < len(file_data):
                    entry, new_offset = _unpack_entry(file_data, offset, current_seq)
                    if entry is None:
                        break  # truncated or corrupt -- stop for this file
                    if entry.sequence >= checkpoint_seq and entry.operation != WALOperation.CHECKPOINT:
                        entries.append(entry)
                    if entry.sequence >= max_seq:
                        max_seq = entry.sequence + 1
                    current_seq += 1
                    offset = new_offset

            # Update sequence to continue from the highest seen
            if max_seq > self._sequence:
                self._sequence = max_seq

        return entries

    # ------------------------------------------------------------------
    # Rotation / Truncation
    # ------------------------------------------------------------------

    def rotate(self) -> None:
        """Rotate to a new WAL file.

        Flushes and closes the current log file, then creates a new one
        with the current sequence number in the filename.

        Raises:
            RuntimeError: If the WAL has been closed.
        """
        with self._lock:
            if self._closed:
                raise RuntimeError("WAL is closed")
            self._rotate_internal()

    def _rotate_internal(self) -> None:
        """Internal rotation -- caller must hold self._lock."""
        if self._current_log is not None:
            self._sync()
            self._current_log.close()
            self._current_log = None
        self._new_log_file()

    def truncate(self, up_to_sequence: int) -> int:
        """Delete WAL files whose entries are all before *up_to_sequence*.

        Args:
            up_to_sequence: Delete files whose entries are entirely below
                            this sequence number.

        Returns:
            Number of files deleted.

        Raises:
            RuntimeError: If the WAL has been closed.
        """
        with self._lock:
            if self._closed:
                raise RuntimeError("WAL is closed")

            wal_files = self._get_wal_files()
            deleted = 0

            for i, wal_file in enumerate(wal_files):
                # Never delete the current log file
                if wal_file == self._current_log_path:
                    continue

                # Check the NEXT file's starting sequence to determine
                # if this file's entries are all below up_to_sequence
                if i + 1 < len(wal_files):
                    next_data = wal_files[i + 1].read_bytes()
                    next_result = _unpack_header(next_data)
                    if next_result is not None:
                        _version, next_seq_start = next_result
                        if next_seq_start <= up_to_sequence:
                            try:
                                wal_file.unlink()
                                deleted += 1
                            except OSError:
                                pass
                else:
                    # Last file before current -- check if all entries
                    # are below up_to_sequence by reading its header
                    file_data = wal_file.read_bytes()
                    result = _unpack_header(file_data)
                    if result is not None:
                        _version, seq_start = result
                        # Count entries to find max sequence
                        current_seq = seq_start
                        offset = WAL_HEADER_SIZE
                        while offset < len(file_data):
                            entry, new_offset = _unpack_entry(file_data, offset, current_seq)
                            if entry is None:
                                break
                            current_seq += 1
                            offset = new_offset
                        if current_seq <= up_to_sequence:
                            try:
                                wal_file.unlink()
                                deleted += 1
                            except OSError:
                                pass

            return deleted

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

    def flush(self) -> None:
        """Flush pending writes to disk without creating a checkpoint.

        Raises:
            RuntimeError: If the WAL has been closed.
        """
        with self._lock:
            if self._closed:
                raise RuntimeError("WAL is closed")
            self._sync()

    def close(self) -> None:
        """Flush and close the WAL file."""
        with self._lock:
            if self._closed:
                return
            self._closed = True
            if self._current_log is not None:
                try:
                    self._current_log.flush()
                    os.fsync(self._current_log.fileno())
                except OSError:
                    pass
                self._current_log.close()
                self._current_log = None

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

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

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

    def __len__(self) -> int:
        """Number of entries since the last checkpoint."""
        return self.entries_since_checkpoint

    def __contains__(self, sequence: int) -> bool:
        """Check if an entry with the given sequence exists in WAL files."""
        with self._lock:
            for wal_file in self._get_wal_files():
                try:
                    file_data = wal_file.read_bytes()
                except OSError:
                    continue
                result = _unpack_header(file_data)
                if result is None:
                    continue
                _version, seq_start = result
                current_seq = seq_start
                offset = WAL_HEADER_SIZE
                while offset < len(file_data):
                    entry, new_offset = _unpack_entry(file_data, offset, current_seq)
                    if entry is None:
                        break
                    if entry.sequence == sequence:
                        return True
                    current_seq += 1
                    offset = new_offset
        return False

    def __repr__(self) -> str:
        return (
            f"WriteAheadLog(log_dir={self._log_dir!r}, "
            f"sequence={self._sequence}, "
            f"entries_since_checkpoint={self._entries_since_checkpoint}, "
            f"log_size={self._current_log_size})"
        )
