#!/usr/bin/env python3
"""Convert pfSense Traffic Totals (vnstat --json) WAN history into OpenMetrics
for `promtool tsdb create-blocks-from openmetrics`.

Uses the finest data vnstat still has: daily totals, then hourly, then 5-minute.
Each period's bytes are spread evenly over 30 s samples of a counter, so
increase() and rate() over the history behave like live data.

The series use the same labels as the live pfSense job from Part 6
(host="pfsense", iface="WAN") plus job="vnstat-import", so every panel that sums
the WAN traffic shows history + live data without changes.

usage: vnstat2prom.py vnstat.json wan-history.om --iface igc0 [--until EPOCH]
  --iface  the WAN interface name as vnstat knows it (pfSense: Interfaces -> Assignments)
  --until  drop history from this Unix time on (when live data starts), so nothing is counted twice
"""
import argparse, json, datetime as dt

ap = argparse.ArgumentParser()
ap.add_argument("src"); ap.add_argument("out")
ap.add_argument("--iface", required=True)
ap.add_argument("--until", type=int, default=None)
a = ap.parse_args()

d = json.load(open(a.src))
ifs = {i["name"]: i for i in d["interfaces"]}
if a.iface not in ifs:
    raise SystemExit("interface %s not in this export; found: %s" % (a.iface, ", ".join(ifs)))
w = ifs[a.iface]["traffic"]

days = sorted(w.get("day", []), key=lambda x: x["timestamp"])
hours = sorted(w.get("hour", []), key=lambda x: x["timestamp"])
fives = sorted(w.get("fiveminute", []), key=lambda x: x["timestamp"])
if not days:
    raise SystemExit("no daily data in this export")

# day data until the first midnight that has complete hourly data, then hours, then 5-minute buckets
periods = []
day_cut = min((x["timestamp"] for x in days if hours and x["timestamp"] >= hours[0]["timestamp"]), default=None)
if day_cut is None:
    day_cut = days[-1]["timestamp"] + 86400
# a day's length comes from the next day's start: 23 or 25 hours across a DST change
for x, nxt in zip(days, days[1:] + [None]):
    end = nxt["timestamp"] if nxt else x["timestamp"] + 86400
    if end <= day_cut:
        periods.append((x["timestamp"], end - x["timestamp"], x["rx"], x["tx"]))
    elif x["timestamp"] < day_cut:
        periods.append((x["timestamp"], day_cut - x["timestamp"], x["rx"], x["tx"]))
five_cut = fives[0]["timestamp"] + (3600 - fives[0]["timestamp"] % 3600) % 3600 if fives else None
for x in hours:
    if x["timestamp"] >= day_cut and (five_cut is None or x["timestamp"] + 3600 <= five_cut):
        periods.append((x["timestamp"], 3600, x["rx"], x["tx"]))
if five_cut is not None:
    end_hours = max([p[0] + p[1] for p in periods] + [0])
    for x in fives:
        if x["timestamp"] >= max(five_cut, end_hours):
            periods.append((x["timestamp"], 300, x["rx"], x["tx"]))
periods.sort()
for p, q in zip(periods, periods[1:]):
    if p[0] + p[1] != q[0]:
        raise SystemExit("gap or overlap in the vnstat data at %d; export again" % p[0])

STEP = 30
LBL = 'device="%s",host="pfsense",iface="WAN",instance="vnstat-import",job="vnstat-import"' % a.iface
rx = tx = 0.0
pts = [(periods[0][0], 0.0, 0.0)]
for start, length, r, t in periods:
    n = max(1, length // STEP)
    for k in range(1, n + 1):
        rx += r / n; tx += t / n
        pts.append((start + k * STEP, rx, tx))
if a.until:
    pts = [p for p in pts if p[0] < a.until]

with open(a.out, "w") as f:
    f.write("# TYPE node_network_receive_bytes counter\n")
    for ts, r, _ in pts:
        f.write("node_network_receive_bytes_total{%s} %.0f %d\n" % (LBL, r, ts))
    f.write("# TYPE node_network_transmit_bytes counter\n")
    for ts, _, t in pts:
        f.write("node_network_transmit_bytes_total{%s} %.0f %d\n" % (LBL, t, ts))
    f.write("# EOF\n")

fmt = lambda t: dt.datetime.fromtimestamp(t).strftime("%Y-%m-%d %H:%M")
print("history %s -> %s, download %.2f GiB, upload %.2f GiB, %d samples" % (
    fmt(pts[0][0]), fmt(pts[-1][0]), pts[-1][1] / 2**30, pts[-1][2] / 2**30, 2 * len(pts)))