import mmap
import os
import struct
import threading
from pathlib import Path

from ..database.writer import (
    pack_record, mfh,
    FILE_HEADER_SIZE, STATUS_LIVE, STATUS_DEAD,
    REC_HEADER_FMT, REC_HEADER_SIZE, CRC_SIZE, crc32,
)
from ..database.reader import scan_all, parse_record_at

CHUNK_SIZE = 256 * 1024


class Engine:
    def __init__(self, path: str):
        self.path = Path(path)
        self._lock = threading.RLock()
        self._index: dict[bytes, tuple[int, int]] = {}
        self._mm = None
        self._f = None
        self._write_pos = FILE_HEADER_SIZE
        self._alloc_size = FILE_HEADER_SIZE
        self._open()

    def _open(self):
        if not self.path.exists():
            with open(self.path, 'wb') as f:
                f.write(mfh())
        self._f = open(self.path, 'r+b')
        self._alloc_size = os.path.getsize(self.path)
        self._build_mmap()
        self._build_index()

    def _build_mmap(self):
        if self._mm and not self._mm.closed:
            self._mm.close()
            self._mm = None
        if self._alloc_size > FILE_HEADER_SIZE:
            self._mm = mmap.mmap(self._f.fileno(), 0)

    def _build_index(self):
        self._index.clear()
        self._write_pos = FILE_HEADER_SIZE
        if not self._mm:
            return
        for offset, rec in scan_all(self._mm):
            if not rec['valid']:
                break
            end = offset + rec['total_size']
            if end > self._write_pos:
                self._write_pos = end
            if rec['status'] == STATUS_LIVE:
                self._index[rec['key']] = (offset, rec['total_size'])
            elif rec['status'] == STATUS_DEAD:
                self._index.pop(rec['key'], None)

    def _ensure_capacity(self, extra: int):
        needed = self._write_pos + extra
        if needed <= self._alloc_size:
            return
        new_alloc = needed + CHUNK_SIZE - (needed % CHUNK_SIZE)
        self._f.seek(new_alloc - 1)
        self._f.write(b'\x00')
        self._f.flush()
        if self._mm and not self._mm.closed:
            try:
                self._mm.resize(new_alloc)
            except (SystemError, OSError):
                self._mm.close()
                self._mm = mmap.mmap(self._f.fileno(), 0)
        else:
            self._mm = mmap.mmap(self._f.fileno(), 0)
        self._alloc_size = new_alloc

    def _close_mm(self):
        if self._mm and not self._mm.closed:
            self._mm.close()
            self._mm = None

    def _ttd(self): # <- truncate_to_data btw
        self._close_mm()
        self._f.flush()
        self._f.truncate(self._write_pos)
        self._alloc_size = self._write_pos
        if self._write_pos > FILE_HEADER_SIZE:
            self._mm = mmap.mmap(self._f.fileno(), 0)

    # ------------------------------------------------------------------
    # Primitives
    # ------------------------------------------------------------------

    def has(self, key: bytes) -> bool:
        return key in self._index

    def read(self, key: bytes):
        with self._lock:
            entry = self._index.get(key)
            if entry is None:
                return None
            offset, _ = entry
            rec = parse_record_at(self._mm, offset)
            if not rec['valid'] or rec['status'] != STATUS_LIVE:
                return None
            return rec['value']

    def append(self, key: bytes, value: bytes) -> int:
        with self._lock:
            packed = pack_record(key, value)
            self._ensure_capacity(len(packed))
            offset = self._write_pos
            self._mm[offset : offset + len(packed)] = packed
            self._write_pos += len(packed)
            self._index[key] = (offset, len(packed))
            return offset

    def append_many(self, pairs: list) -> None:
        with self._lock:
            packed_list = [(k, v, pack_record(k, v)) for k, v in pairs]
            total = sum(len(p) for _, _, p in packed_list)
            self._ensure_capacity(total)
            for key, _, packed in packed_list:
                offset = self._write_pos
                self._mm[offset : offset + len(packed)] = packed
                self._write_pos += len(packed)
                self._index[key] = (offset, len(packed))

    def _tombstone_at(self, offset: int, rec: dict):
        """
        Flip the status byte to DEAD and recompute CRC in-place.
        CRC must be updated — otherwise the scan on reopen will see a CRC mismatch and stop, 
        losing all records that follow. So, we need to be careful with this.
        """
        self._mm[offset : offset + 1] = bytes([STATUS_DEAD])
        payload_end = offset + REC_HEADER_SIZE + rec['key_len'] + rec['val_len']
        payload = bytes(self._mm[offset : payload_end])
        new_crc = struct.pack('>I', crc32(payload))
        self._mm[payload_end : payload_end + CRC_SIZE] = new_crc

    def tombstone(self, key: bytes) -> bool:
        with self._lock:
            entry = self._index.get(key)
            if entry is None:
                return False
            offset, _ = entry
            rec = parse_record_at(self._mm, offset)
            self._tombstone_at(offset, rec)
            del self._index[key]
            return True

    def tombstone_many(self, keys: list) -> int:
        with self._lock:
            count = 0
            for key in keys:
                entry = self._index.get(key)
                if entry is None:
                    continue
                offset, _ = entry
                rec = parse_record_at(self._mm, offset)
                self._tombstone_at(offset, rec)
                del self._index[key]
                count += 1
            return count

    def overwrite_value(self, key: bytes, value: bytes) -> bool:
        with self._lock:
            entry = self._index.get(key)
            if entry is None:
                return False
            offset, _ = entry
            rec = parse_record_at(self._mm, offset)
            if len(value) != rec['val_len']:
                return False
            val_start = offset + REC_HEADER_SIZE + rec['key_len']
            val_end   = val_start + rec['val_len']
            self._mm[val_start : val_end] = value
            payload  = bytes(self._mm[offset : val_end])
            new_crc  = struct.pack('>I', crc32(payload))
            self._mm[val_end : val_end + CRC_SIZE] = new_crc
            return True

    def flush(self):
        with self._lock:
            if self._mm and not self._mm.closed:
                self._mm.flush()

    # ------------------------------------------------------------------
    # Compaction
    # ------------------------------------------------------------------

    def compact(self):
        with self._lock:
            tmp = self.path.with_suffix('.lxdb.tmp')
            try:
                with open(tmp, 'wb') as out:
                    out.write(mfh())
                    new_index     = {}
                    new_write_pos = FILE_HEADER_SIZE
                    src_mm = mmap.mmap(self._f.fileno(), 0, access=mmap.ACCESS_READ)
                    try:
                        for offset, rec in scan_all(src_mm):
                            if rec['valid'] and rec['status'] == STATUS_LIVE:
                                packed  = pack_record(rec['key'], rec['value'])
                                new_off = out.tell()
                                out.write(packed)
                                new_index[rec['key']] = (new_off, len(packed))
                                new_write_pos = new_off + len(packed)
                    finally:
                        src_mm.close()
                self._close_mm()
                self._f.close()
                os.replace(tmp, self.path)
                self._f          = open(self.path, 'r+b')
                self._alloc_size = os.path.getsize(self.path)
                self._write_pos  = new_write_pos
                if self._alloc_size > FILE_HEADER_SIZE:
                    self._mm = mmap.mmap(self._f.fileno(), 0)
                self._index = new_index
            except Exception:
                if tmp.exists():
                    tmp.unlink()
                raise

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

    def close(self):
        with self._lock:
            if self._mm and not self._mm.closed:
                self._mm.flush()
            if self._write_pos < self._alloc_size:
                self._ttd()
            else:
                self._close_mm()
            if self._f:
                self._f.close()
                self._f = None

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()