"""Share identical blocks between files on file systems that support it (btrfs, XFS): the kernel compares the bytes
itself and only then lets the two files point at the same disk blocks (FIDEDUPERANGE), so contents never change.
Serve caches of two eras are mostly the same bytes at the same offsets (only the areas an era changes differ):
sharing them roughly halves the disk they take. Best effort: anywhere else it does nothing."""
from __future__ import annotations

import fcntl
import os
import struct
from pathlib import Path

FIDEDUPERANGE = 0xC0189436          # _IOWR(0x94, 54, struct file_dedupe_range)
BLOCK = 128 << 10                   # compare in btrfs compressed-extent units
MAX_CALL = 16 << 20                 # btrfs dedupes at most 16 MiB per call
READ = 16 << 20


def _ioctl(src_fd: int, dst_fd: int, off: int, length: int) -> int:
    """Ask the kernel to share [off, off+length) of dst with the same range of src; returns bytes shared."""
    buf = bytearray(struct.pack("<QQHHI", off, length, 1, 0, 0) + struct.pack("<qQQiI", dst_fd, off, 0, 0, 0))
    fcntl.ioctl(src_fd, FIDEDUPERANGE, buf)
    _, _, done, status, _ = struct.unpack_from("<qQQiI", buf, 24)
    return done if status == 0 else 0


def same_runs(a, b, size: int, block: int = BLOCK):
    """(offset, length) runs of whole blocks that are byte-identical in the two open files."""
    runs, start, off = [], None, 0
    while off + block <= size:
        n = min(READ, (size - off) // block * block)
        x, y = os.pread(a, n, off), os.pread(b, n, off)
        for k in range(0, n, block):
            if x[k:k + block] == y[k:k + block]:
                if start is None:
                    start = off + k
            elif start is not None:
                runs.append((start, off + k - start))
                start = None
        off += n
    if start is not None:
        runs.append((start, off - start))
    return runs


def dedupe_file(src: Path, dst: Path) -> int:
    """Share dst's blocks that equal src's (same offsets). Bytes shared; 0 if unsupported or nothing matches."""
    size = os.path.getsize(src)
    if size != os.path.getsize(dst) or size < BLOCK:
        return 0
    a, b = os.open(src, os.O_RDONLY), os.open(dst, os.O_RDONLY)
    try:
        shared = 0
        for off, length in same_runs(a, b, size):
            for o in range(off, off + length, MAX_CALL):
                shared += _ioctl(a, b, o, min(MAX_CALL, off + length - o))
        return shared
    except OSError:                                    # tmpfs, ext4, permissions: nothing to share
        return 0
    finally:
        os.close(a)
        os.close(b)


def dedupe_dirs(new: Path, peers) -> int:
    """Share new's files with same-named, same-sized files in peer folders. Bytes shared."""
    shared = 0
    for peer in peers:
        for f in sorted(Path(new).iterdir()):
            p = Path(peer) / f.name
            try:
                if f.suffix == ".npy" and p.is_file():
                    shared += dedupe_file(p, f)
            except OSError:                            # a peer pruned meanwhile: skip it
                pass
    return shared