#!/usr/bin/env python3
"""Write amplification in an LSM-tree, DDIA 2e Ch. 4.

One logical write turns into many physical writes. In an LSM-tree, a value is
written to the write-ahead log for durability, again when the memtable is
flushed to an SSTable, and AGAIN every time compaction rewrites the SSTable it
lives in as it migrates down the levels. The book defines write amplification as
the total bytes written to disk divided by the bytes a plain append-only log
would have written.

We insert ~230 MB of key/value data into a real LSM engine (RocksDB) with a
small memtable so the data cascades through several levels, then read RocksDB's
own byte counters to compute the write amplification. Values are random
(incompressible) so on-disk SSTable size tracks the logical data, and the
rewrites are visible rather than hidden by compression.

Run it:
    uv run --with rocksdict python3 write_amplification.py
(Needs RocksDB via the `rocksdict` package. The exact factor varies a little
between runs -- background compaction is timing-dependent -- but the shape holds:
a single-digit multiple, dominated by compaction.)
"""
import random
import shutil
import tempfile

from rocksdict import Rdict, Options

try:
    from rocksdict import DBCompressionType
    _HAVE_CT = True
except Exception:
    _HAVE_CT = False

N = 2_000_000
VALUE_LEN = 100


def ticker_bytes(stats_text):
    st = {}
    for line in stats_text.splitlines():
        parts = line.split(" COUNT : ")
        if len(parts) == 2 and parts[1].strip().isdigit():
            st[parts[0].strip()] = int(parts[1].strip())
    return st


def main():
    path = tempfile.mkdtemp(prefix="ddia_wa_")
    opt = Options(raw_mode=True)
    opt.create_if_missing(True)
    opt.enable_statistics()
    opt.set_write_buffer_size(4 * 1024 * 1024)          # 4 MB memtable -> frequent flushes
    opt.set_max_bytes_for_level_base(8 * 1024 * 1024)   # small levels -> data cascades down
    opt.set_max_bytes_for_level_multiplier(4)
    if _HAVE_CT:
        try:
            opt.set_compression_type(DBCompressionType.none())
        except Exception:
            pass

    db = Rdict(path, opt)
    rnd = random.Random(42)
    user_bytes = 0
    for i in range(N):
        key = f"key{(i * 2654435761) % (2 ** 32):012d}".encode()   # unique, scattered
        value = rnd.randbytes(VALUE_LEN)                           # incompressible
        db[key] = value
        user_bytes += len(key) + len(value)

    db.flush()
    db.compact_range(None, None)   # settle all pending compaction

    st = ticker_bytes(opt.get_statistics())
    wal = st.get("rocksdb.wal.bytes", 0)
    flush = st.get("rocksdb.flush.write.bytes", 0)
    compact = st.get("rocksdb.compact.write.bytes", 0)
    total = wal + flush + compact
    mb = 1e6

    print(f"inserted {N:,} records (~{user_bytes / mb:.0f} MB of key+value data) "
          f"into RocksDB (an LSM engine)\n")
    print(f"what RocksDB actually wrote to disk to store {user_bytes / mb:.0f} MB durably:")
    print(f"    write-ahead log (WAL, written once)   = {wal / mb:6.0f} MB")
    print(f"    memtable flushes to L0                = {flush / mb:6.0f} MB")
    print(f"    compaction (rewriting SSTables)       = {compact / mb:6.0f} MB")
    print(f"    ----------------------------------------------------")
    print(f"    total written to disk                 = {total / mb:6.0f} MB\n")
    print(f"write amplification = {total / mb:.0f} / {user_bytes / mb:.0f} "
          f"= {total / user_bytes:.1f}x   (vs a plain append-only log)")
    print(f"  compaction alone rewrote the data {compact / user_bytes:.1f}x "
          f"as it migrated down the levels")

    final = 0
    try:
        final = db.property_int_value("rocksdb.total-sst-files-size") or 0
    except Exception:
        pass
    if final:
        print(f"\nfinal database size on disk = {final / mb:.0f} MB "
              f"(~= the {user_bytes / mb:.0f} MB of data; the rewrites are transient work, not waste)")

    db.close()
    Rdict.destroy(path, opt)
    shutil.rmtree(path, ignore_errors=True)


if __name__ == "__main__":
    main()