#!/usr/bin/env python3
"""
Work-Day Reliability logger for Starlink Remote Work USA.

Runs on an always-on box (Raspberry Pi / mini PC) wired to the owner's own Starlink connection and writes:

  {out}/{yyyy-mm}.csv          one row per minute
      ts, latency_ms, jitter_ms, loss_pct, outage_s, dish_latency_ms, dish_drop_pct, obstructed_pct,
      down_mbps, up_mbps, test
  {out}/{yyyy-mm}-outages.csv  one row per outage (>= --outage-min seconds with no reply from ANY target)
      start, end, seconds

How each column is measured
  latency_ms   median round-trip time of all replies to the probe targets during the minute's probe burst
  jitter_ms    mean absolute difference between consecutive RTTs (per target, averaged) — RFC 3550-style
  loss_pct     % of probes in the burst that got no reply
  outage_s     seconds in this minute flagged "down" by the 1-second heartbeat (no reply from any target)
  dish_*       Starlink dish gRPC status (pop ping latency, pop ping drop rate, fraction obstructed), via
               starlink-grpc-tools if installed; blank otherwise. Read-only: the logger never changes settings.
  down/up_mbps hourly iperf3 test to your own server (--iperf HOST); blank when no test ran that minute
  test         "speed" when a speed test ran in that minute

Nothing here is sent anywhere. The CSVs stay on the box until the owner copies them into the site's
src/data/report/ folder. Python 3.9+, standard library only (starlink-grpc-tools optional).

Usage
  python3 workday_logger.py --out ~/workday-data --targets 1.1.1.1 8.8.8.8 9.9.9.9 [--iperf my.vps.example]
  python3 workday_logger.py --self-test        # checks the parsing and math, no network
  python3 workday_logger.py --once --out /tmp/x # one minute, then exit (to try it)
"""
from __future__ import annotations

import argparse
import csv
import datetime as dt
import json
import os
import platform
import re
import shutil
import statistics
import subprocess
import sys
import threading
import time

MINUTE_COLUMNS = ["ts", "latency_ms", "jitter_ms", "loss_pct", "outage_s", "dish_latency_ms", "dish_drop_pct",
                  "obstructed_pct", "down_mbps", "up_mbps", "test"]
OUTAGE_COLUMNS = ["start", "end", "seconds"]
RTT_RE = re.compile(r"time[=<]([\d.]+)\s*ms")
IS_MAC = platform.system() == "Darwin"


# ---------------------------------------------------------------- pure helpers (self-tested)

def parse_ping_rtts(output: str) -> list[float]:
    """RTTs (ms) from `ping` output on Linux or macOS."""
    return [float(m) for m in RTT_RE.findall(output)]


def jitter(rtts: list[float]) -> float | None:
    """Mean absolute difference between consecutive RTTs; None with fewer than 2 replies."""
    if len(rtts) < 2:
        return None
    return sum(abs(b - a) for a, b in zip(rtts, rtts[1:])) / (len(rtts) - 1)


def summarize_burst(per_target: dict[str, list[float]], sent_per_target: int) -> dict:
    replies = [r for rtts in per_target.values() for r in rtts]
    sent = sent_per_target * max(1, len(per_target))
    jitters = [j for j in (jitter(r) for r in per_target.values()) if j is not None]
    return {
        "latency_ms": round(statistics.median(replies), 1) if replies else "",
        "jitter_ms": round(sum(jitters) / len(jitters), 1) if jitters else "",
        "loss_pct": round(100 * (sent - len(replies)) / sent, 1) if sent else "",
    }


def local_iso(t: float | None = None) -> str:
    """Local wall-clock time with UTC offset, e.g. 2026-11-02T08:15:00-06:00 (seconds kept)."""
    return dt.datetime.fromtimestamp(t if t is not None else time.time()).astimezone().replace(microsecond=0).isoformat()


# ---------------------------------------------------------------- probes

def ping_cmd(target: str, count: int, interval: float, timeout_s: int) -> list[str]:
    if IS_MAC:  # macOS: -W is in milliseconds; -t is the overall timeout
        return ["ping", "-n", "-c", str(count), "-i", str(interval), "-W", str(timeout_s * 1000), target]
    return ["ping", "-n", "-c", str(count), "-i", str(interval), "-W", str(timeout_s), target]


def ping_burst(target: str, count: int = 20, interval: float = 0.25) -> list[float]:
    try:
        out = subprocess.run(ping_cmd(target, count, interval, 1), capture_output=True, text=True,
                             timeout=count * interval + 5).stdout
    except (subprocess.TimeoutExpired, OSError):
        return []
    return parse_ping_rtts(out)


def ping_once(target: str) -> bool:
    try:
        r = subprocess.run(ping_cmd(target, 1, 1, 1), capture_output=True, text=True, timeout=3)
        return r.returncode == 0
    except (subprocess.TimeoutExpired, OSError):
        return False


def dish_status() -> dict:
    """Read-only dish stats via starlink-grpc-tools (pip install starlink-grpc-tools). Blank if unavailable."""
    blank = {"dish_latency_ms": "", "dish_drop_pct": "", "obstructed_pct": ""}
    try:
        import starlink_grpc  # type: ignore
    except ImportError:
        return blank
    try:
        status, obstruction, _alerts = starlink_grpc.status_data()
    except Exception:  # dish unreachable, API changed: never crash the logger over it
        return blank

    def pick(key, scale=1.0):
        v = status.get(key, obstruction.get(key) if isinstance(obstruction, dict) else None)
        return "" if v is None else round(float(v) * scale, 3)

    return {"dish_latency_ms": pick("pop_ping_latency_ms"), "dish_drop_pct": pick("pop_ping_drop_rate", 100),
            "obstructed_pct": pick("fraction_obstructed", 100)}


def iperf(host: str, seconds: int = 10) -> tuple[str, str]:
    """(down_mbps, up_mbps) from iperf3 to the owner's own server; blanks on failure."""
    if not shutil.which("iperf3"):
        return "", ""

    def run(reverse: bool) -> str:
        args = ["iperf3", "-c", host, "-t", str(seconds), "-J"] + (["-R"] if reverse else [])
        try:
            out = subprocess.run(args, capture_output=True, text=True, timeout=seconds + 20).stdout
            bps = json.loads(out)["end"]["sum_received"]["bits_per_second"]
            return str(round(bps / 1e6, 1))
        except Exception:
            return ""

    return run(True), run(False)


# ---------------------------------------------------------------- heartbeat (outages)

class Heartbeat(threading.Thread):
    """Pings the targets once a second; a second is 'down' when none of them answers."""

    def __init__(self, targets: list[str], out_dir: str, outage_min: int):
        super().__init__(daemon=True)
        self.targets, self.out_dir, self.outage_min = targets, out_dir, outage_min
        self.lock = threading.Lock()
        self.down_by_minute: dict[str, int] = {}
        self.down_start: float | None = None

    def run(self):
        while True:
            t0 = time.time()
            up = any(ping_once(t) for t in self.targets[:2])
            key = local_iso(t0)[:16]
            with self.lock:
                if not up:
                    self.down_by_minute[key] = self.down_by_minute.get(key, 0) + 1
                    if self.down_start is None:
                        self.down_start = t0
                elif self.down_start is not None:
                    self._close(t0)
            time.sleep(max(0.0, 1.0 - (time.time() - t0)))

    def _close(self, t_end: float):
        seconds = round(t_end - self.down_start)
        if seconds >= self.outage_min:
            path = os.path.join(self.out_dir, f"{local_iso(self.down_start)[:7]}-outages.csv")
            append_row(path, OUTAGE_COLUMNS, {"start": local_iso(self.down_start), "end": local_iso(t_end), "seconds": seconds})
        self.down_start = None

    def take(self, minute_key: str) -> int:
        with self.lock:
            return min(60, self.down_by_minute.pop(minute_key, 0))


# ---------------------------------------------------------------- output

def append_row(path: str, columns: list[str], row: dict):
    new = not os.path.exists(path)
    with open(path, "a", newline="") as f:
        w = csv.DictWriter(f, fieldnames=columns, extrasaction="ignore")
        if new:
            w.writeheader()
        w.writerow({c: row.get(c, "") for c in columns})


def one_minute(args, hb: Heartbeat | None, minute_start: float) -> dict:
    results: dict[str, list[float]] = {}
    threads = [threading.Thread(target=lambda t=t: results.__setitem__(t, ping_burst(t, args.count, args.interval)))
               for t in args.targets]
    for th in threads:
        th.start()
    for th in threads:
        th.join()
    row = {"ts": local_iso(minute_start), **summarize_burst(results, args.count), **dish_status()}
    lt = time.localtime(minute_start)
    if args.iperf and lt.tm_min == args.speed_minute:
        row["down_mbps"], row["up_mbps"] = iperf(args.iperf)
        row["test"] = "speed"
    return row


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--out", default=os.path.expanduser("~/workday-data"))
    p.add_argument("--targets", nargs="+", default=["1.1.1.1", "8.8.8.8", "9.9.9.9"])
    p.add_argument("--count", type=int, default=20, help="probes per target per minute")
    p.add_argument("--interval", type=float, default=0.25, help="seconds between probes (>=0.2 without root)")
    p.add_argument("--iperf", help="your own iperf3 server for the hourly speed test")
    p.add_argument("--speed-minute", type=int, default=30, help="minute of the hour the speed test runs")
    p.add_argument("--outage-min", type=int, default=2, help="shortest gap (s) written to the outages file")
    p.add_argument("--once", action="store_true")
    p.add_argument("--self-test", action="store_true")
    args = p.parse_args()

    if args.self_test:
        return self_test()
    os.makedirs(args.out, exist_ok=True)
    if args.once:
        row = one_minute(args, None, time.time())
        row["outage_s"] = 0
        append_row(os.path.join(args.out, f"{row['ts'][:7]}.csv"), MINUTE_COLUMNS, row)
        print(json.dumps(row), flush=True)
        return 0
    hb = Heartbeat(args.targets, args.out, args.outage_min)
    hb.start()
    pending = None  # the row for the minute in progress; written once the minute is over
    while True:
        time.sleep(60 - (time.time() % 60))
        now = time.time()
        if pending:  # the minute just ended: its heartbeat count is now complete
            pending["outage_s"] = hb.take(pending["ts"][:16])
            append_row(os.path.join(args.out, f"{pending['ts'][:7]}.csv"), MINUTE_COLUMNS, pending)
            print(json.dumps(pending), flush=True)
        pending = one_minute(args, hb, now - (now % 60))  # probe burst first, then any speed test


def self_test():
    linux = "64 bytes from 1.1.1.1: icmp_seq=1 ttl=57 time=31.4 ms\n64 bytes from 1.1.1.1: icmp_seq=2 ttl=57 time=35.0 ms\n"
    mac = "64 bytes from 1.1.1.1: icmp_seq=0 ttl=57 time=28.112 ms\n"
    assert parse_ping_rtts(linux) == [31.4, 35.0], parse_ping_rtts(linux)
    assert parse_ping_rtts(mac) == [28.112]
    assert parse_ping_rtts("Request timeout for icmp_seq 0") == []
    assert jitter([30, 34, 31]) == 3.5
    assert jitter([30]) is None
    s = summarize_burst({"a": [30, 32, 34], "b": [40]}, 4)
    assert s["latency_ms"] == 33.0 and s["loss_pct"] == 50.0 and s["jitter_ms"] == 2.0, s
    s = summarize_burst({"a": [], "b": []}, 4)
    assert s["latency_ms"] == "" and s["loss_pct"] == 100.0
    assert re.match(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}[+-]\d{2}:\d{2}$", local_iso(0)), local_iso(0)
    print("self-test ok")
    return 0


if __name__ == "__main__":
    sys.exit(main())