"""Tiny Prometheus exporter for scheduled internet speed tests (Ookla CLI).

Runs a test at the hours listed in SPEEDTEST_HOURS (local time, TZ env) and
serves the last result on :9798/metrics. Prometheus only reads the stored
result, so scraping never starts a test.
"""
import json, os, subprocess, threading, time, datetime
from prometheus_client import Gauge, Counter, start_http_server

HOURS = sorted({int(h) for h in os.environ.get("SPEEDTEST_HOURS", "8,14,20").split(",") if h.strip()})
SERVER = os.environ.get("SPEEDTEST_SERVER", "").strip()
TIMEOUT = int(os.environ.get("SPEEDTEST_TIMEOUT", "120"))
PORT = int(os.environ.get("SPEEDTEST_PORT", "9798"))

down = Gauge("speedtest_download_bits_per_second", "Download speed of the last test")
up = Gauge("speedtest_upload_bits_per_second", "Upload speed of the last test")
ping = Gauge("speedtest_ping_latency_milliseconds", "Idle latency to the test server")
jitter = Gauge("speedtest_jitter_latency_milliseconds", "Jitter to the test server")
loss = Gauge("speedtest_packet_loss_percent", "Packet loss during the test (if reported)")
used = Gauge("speedtest_data_used_bytes", "Data transferred by the last test (down + up)")
last_ok = Gauge("speedtest_last_success_timestamp_seconds", "Unix time of the last successful test")
last_run = Gauge("speedtest_last_run_timestamp_seconds", "Unix time of the last attempt")
success = Gauge("speedtest_success", "1 if the last attempt worked, 0 if it failed")
runs = Counter("speedtest_runs", "Speed tests attempted", ["result"])


def run_test():
    cmd = ["speedtest", "--accept-license", "--accept-gdpr", "-f", "json"]
    if SERVER:
        cmd += ["-s", SERVER]
    last_run.set(time.time())
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=TIMEOUT, check=True).stdout
        r = json.loads(out.strip().splitlines()[-1])
        down.set(r["download"]["bandwidth"] * 8)
        up.set(r["upload"]["bandwidth"] * 8)
        ping.set(r["ping"]["latency"])
        jitter.set(r["ping"].get("jitter", 0))
        if r.get("packetLoss") is not None:
            loss.set(r["packetLoss"])
        used.set(r["download"].get("bytes", 0) + r["upload"].get("bytes", 0))
        last_ok.set(time.time())
        success.set(1)
        runs.labels("ok").inc()
        print("%s ok: %.0f / %.0f Mbit/s, ping %.1f ms, server %s" % (
            datetime.datetime.now().strftime("%F %T"), r["download"]["bandwidth"] * 8 / 1e6,
            r["upload"]["bandwidth"] * 8 / 1e6, r["ping"]["latency"], r.get("server", {}).get("name", "?")), flush=True)
    except Exception as e:  # keep the exporter alive whatever happens
        success.set(0)
        runs.labels("failed").inc()
        print("%s FAILED: %s" % (datetime.datetime.now().strftime("%F %T"), e), flush=True)


def scheduler():
    if os.environ.get("SPEEDTEST_ON_START", "1") == "1":
        run_test()
    done = None
    while True:
        now = datetime.datetime.now()
        key = (now.date(), now.hour)
        if now.hour in HOURS and now.minute < 10 and key != done:
            done = key
            run_test()
        time.sleep(30)


if __name__ == "__main__":
    start_http_server(PORT)
    print("speedtest exporter on :%d, tests at %s:00 (%s)" % (PORT, ",".join(map(str, HOURS)), os.environ.get("TZ", "UTC")), flush=True)
    scheduler()