Godstrument/recorder.py

"""
recorder.py — the instrument's memory. Live and historical in one file.

Taps the hub's WebSocket firehose and writes every signal + event to a local
SQLite database (`godstrument.db`). Once it's recording, ANY source — including
your own hand, the room mic, the ToF sensor — can be replayed and time-warped
later by workers/replay_db.py, not just the feeds that happen to have a public
history API.

It's cheap: change-based logging (only write when a signal actually moves) plus a
decimation cap and a slow heartbeat. Slow world data costs almost nothing; only
fast local sensors add up, and even those are bounded.

    python recorder.py                 # record whatever the hub is broadcasting
    python recorder.py --stats         # show size / signals / span, then exit
    python recorder.py --decimate 20   # cap any one signal at 20 Hz on disk
"""

from __future__ import annotations
import argparse
import asyncio
import json
import os
import sqlite3
import time

import websockets

HERE = os.path.dirname(os.path.abspath(__file__))

SCHEMA = """
CREATE TABLE IF NOT EXISTS samples (signal TEXT, t REAL, v REAL);
CREATE INDEX IF NOT EXISTS idx_samples ON samples(signal, t);
CREATE TABLE IF NOT EXISTS events  (signal TEXT, t REAL, mag REAL, meta TEXT);
CREATE INDEX IF NOT EXISTS idx_events ON events(signal, t);
"""


def open_db(path: str) -> sqlite3.Connection:
    con = sqlite3.connect(path)
    con.execute("PRAGMA journal_mode=WAL")
    con.executescript(SCHEMA)
    return con


def stats(path: str):
    if not os.path.exists(path):
        print(f"no database at {path} yet.")
        return
    con = open_db(path)
    con.execute("PRAGMA wal_checkpoint(TRUNCATE)")   # fold WAL in for a true size
    con.commit()
    size = os.path.getsize(path)
    ns = con.execute("SELECT COUNT(*) FROM samples").fetchone()[0]
    ne = con.execute("SELECT COUNT(*) FROM events").fetchone()[0]
    span = con.execute("SELECT MIN(t), MAX(t) FROM samples").fetchone()
    sigs = con.execute("SELECT signal, COUNT(*) c FROM samples "
                       "GROUP BY signal ORDER BY c DESC").fetchall()
    dur = (span[1] - span[0]) if span[0] else 0
    print(f"db: {path}")
    print(f"size: {size/1e6:.2f} MB   samples: {ns:,}   events: {ne:,}")
    if dur:
        print(f"span: {dur/3600:.2f} h   ->  {size/max(dur,1)*86400/1e6:.1f} MB/day"
              f"   ({size/max(ns,1):.1f} bytes/sample)")
    print("busiest signals:")
    for sig, c in sigs[:8]:
        print(f"   {sig:22} {c:,}")


async def record(url, path, decimate, heartbeat, eps):
    con = open_db(path)
    min_iv = 1.0 / decimate if decimate > 0 else 0.0
    last, last_t, seen = {}, {}, set()
    sbuf, ebuf, last_commit = [], [], time.time()
    print(f"[recorder] -> {path}  (decimate {decimate}Hz, heartbeat {heartbeat}s)")

    while True:
        try:
            async with websockets.connect(url, max_size=None) as ws:
                async for msg in ws:
                    d = json.loads(msg)
                    if "sources" not in d:
                        continue
                    now = time.time()
                    for sig, s in d["sources"].items():
                        raw = s.get("raw")
                        if raw is None:
                            continue
                        lv, lt = last.get(sig), last_t.get(sig, 0.0)
                        moved = lv is None or abs(raw - lv) > eps * max(1.0, abs(lv))
                        if (moved and now - lt >= min_iv) or now - lt >= heartbeat:
                            sbuf.append((sig, now, float(raw)))
                            last[sig], last_t[sig] = raw, now
                    for e in d.get("events", []):
                        k = (e.get("src"), round(e.get("t", 0), 3))
                        if k not in seen:
                            seen.add(k)
                            ebuf.append((e.get("src"), now, float(e.get("mag", 0)), None))
                            if len(seen) > 8000:
                                seen = set(list(seen)[-3000:])
                    if now - last_commit > 1.5:
                        if sbuf:
                            con.executemany("INSERT INTO samples VALUES (?,?,?)", sbuf)
                            sbuf = []
                        if ebuf:
                            con.executemany("INSERT INTO events VALUES (?,?,?,?)", ebuf)
                            ebuf = []
                        con.commit()
                        last_commit = now
        except Exception as e:  # noqa
            print(f"[recorder] disconnected ({e}); retrying in 2s")
            await asyncio.sleep(2)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="ws://127.0.0.1:8765")
    ap.add_argument("--db", default=os.path.join(HERE, "godstrument.db"))
    ap.add_argument("--decimate", type=float, default=20.0, help="max Hz per signal")
    ap.add_argument("--heartbeat", type=float, default=10.0, help="force-write every Ns")
    ap.add_argument("--eps", type=float, default=0.0005, help="relative change to log")
    ap.add_argument("--stats", action="store_true")
    args = ap.parse_args()

    if args.stats:
        stats(args.db)
        return
    try:
        asyncio.run(record(args.url, args.db, args.decimate,
                           args.heartbeat, args.eps))
    except KeyboardInterrupt:
        print("\n[recorder] stopped.")


if __name__ == "__main__":
    main()