#!/usr/bin/env python3
"""Measure what SQLite's durability settings cost on a home server.

Every self-hosted app that stores its state in SQLite - Home Assistant,
Immich, Miniflux, Paperless, Gitea in its default setup - inherits whatever
journal_mode and synchronous the app happened to set. This measures the write
throughput of each combination, plus the effect of committing per row versus
per batch, so the durability trade-off can be priced instead of guessed.

It measures speed only. It does NOT test crash safety: pulling power from a
machine is the only honest test of that, and the semantics are documented in
SQLite's own write-ahead log and pragma pages.

Usage: ./sqlite-durability-bench.py [rows] > sqlite-results.jsonl
"""
import json, os, sqlite3, sys, tempfile, time, platform, subprocess

ROWS = int(sys.argv[1]) if len(sys.argv) > 1 else 20_000
COMBOS = [(j, s) for j in ("delete", "wal") for s in ("full", "normal", "off")]
# macOS is the interesting case: plain fsync() there returns once the data has
# reached the drive, not once the drive has committed it to stable storage.
# PRAGMA fullfsync=ON asks for F_FULLFSYNC instead, which is what synchronous=FULL
# is usually assumed to already be doing. It is off by default. Measured both ways.
FULLFSYNC = (False, True)
BATCHES = [1, 100, 1000]
# One row per transaction with F_FULLFSYNC costs about 10 ms per commit, so the
# per-row cases run a smaller number of rows. Throughput is per second either
# way, and every record states the row count it was measured over.
ROWS_FOR = lambda batch: 2_000 if batch == 1 else None  # None means "use ROWS"

DDL = """CREATE TABLE reading (
  id INTEGER PRIMARY KEY, sensor TEXT NOT NULL, ts INTEGER NOT NULL,
  value REAL NOT NULL, unit TEXT NOT NULL);
CREATE INDEX ix_reading_sensor_ts ON reading(sensor, ts);"""

# A sensor-history shaped row, which is what these apps actually write.
ROW = ("balcony-temp", 1_760_000_000, 21.375, "C")


def measure(journal, sync, batch, rows, fullfsync=False):
    path = os.path.join(tempfile.mkdtemp(), "bench.sqlite")
    db = sqlite3.connect(path, isolation_level=None)
    db.execute(f"PRAGMA journal_mode={journal}")
    db.execute(f"PRAGMA synchronous={sync}")
    db.execute(f"PRAGMA fullfsync={'ON' if fullfsync else 'OFF'}")
    db.executescript(DDL)
    sql = "INSERT INTO reading(sensor, ts, value, unit) VALUES (?,?,?,?)"
    t = time.monotonic()
    for i in range(0, rows, batch):
        n = min(batch, rows - i)
        db.execute("BEGIN")
        db.executemany(sql, [(ROW[0], ROW[1] + i + k, ROW[2], ROW[3]) for k in range(n)])
        db.execute("COMMIT")
    elapsed = time.monotonic() - t
    db.execute("PRAGMA wal_checkpoint(TRUNCATE)")
    on_disk = sum(os.path.getsize(path + s) for s in ("", "-wal", "-journal")
                  if os.path.exists(path + s))
    assert db.execute("SELECT count(*) FROM reading").fetchone()[0] == rows
    db.close()
    for s in ("", "-wal", "-journal"):
        if os.path.exists(path + s):
            os.remove(path + s)
    os.rmdir(os.path.dirname(path))
    return elapsed, on_disk


def main():
    cpu = subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"],
                         capture_output=True, text=True).stdout.strip()
    print(json.dumps({"record": "machine", "cpu": cpu, "cores": os.cpu_count(),
                      "os": f"{platform.system()} {platform.release()}",
                      "sqlite": sqlite3.sqlite_version, "rows": ROWS}), flush=True)
    for journal, sync in COMBOS:
        for batch in BATCHES:
          for ff in FULLFSYNC:
            rows = ROWS_FOR(batch) or ROWS
            elapsed, on_disk = measure(journal, sync, batch, rows, ff)
            print(json.dumps({
                "record": "measure", "journal_mode": journal, "synchronous": sync,
                "fullfsync": ff, "rows_per_transaction": batch, "rows": rows,
                "elapsed_s": round(elapsed, 3),
                "rows_per_s": round(rows / elapsed),
                "db_bytes": on_disk,
            }), flush=True)


if __name__ == "__main__":
    main()