#!/usr/bin/python3 -I
import json
import os
import statistics
import time
import urllib.request

API = os.environ.get("DNSDIST_API", "http://127.0.0.1:8083/api/v1/servers/localhost")
KEY = open(os.path.join(os.environ.get("CREDENTIALS_DIRECTORY", "/etc/dnsbunker-stats"), "apikey")).read().strip()
RUN = os.environ.get("RUNTIME_DIRECTORY", "/run/dnsbunker-stats")
STATE, OUT = os.path.join(RUN, "state.json"), os.path.join(RUN, "stats.json")
WINDOW = 86400
STEP = 600
AVAIL_STEP = 1800
LISTS = {
    "tif": (os.environ.get("RPZ_TIF", "/etc/powerdns/rpz/40-tif.rpz"), "hagezi-tif"),
    "pro": (os.environ.get("RPZ_PRO", "/etc/powerdns/rpz/30-pro.rpz"), "hagezi-pro"),
}
CHECK_STAMP = os.environ.get("RPZ_CHECK_STAMP", "/var/lib/systemd/timers/stamp-rpz-update.timer")
REC_API = os.environ.get("REC_API", "")

def get(url, key=None):
    req = urllib.request.Request(url, headers={"X-API-Key": key or KEY})
    with urllib.request.urlopen(req, timeout=10) as r:
        return json.load(r)

def qps_median(samples):
    rates = [s[1] / s[3] for s in samples if len(s) > 3 and s[3] > 0]
    return statistics.median(rates) if rates else 0

def series(samples, now):
    valid = [s for s in samples if len(s) > 3 and s[3] > 0]
    if not valid:
        return [], STEP
    step = min(STEP, max(60, int(now - valid[0][0]) // 48 // 60 * 60))
    buckets = {}
    for s in valid:
        b = buckets.setdefault(int((now - s[0]) // step), [0, 0.0])
        b[0] += s[1]
        b[1] += s[3]
    out, prev = [], None
    for i in range(max(buckets), -1, -1):
        b = buckets.get(i)
        if b and b[1] > 0:
            prev = round(b[0] / b[1], 1)
        if prev is not None:
            out.append(prev)
    return out, step

def availability(samples, now):
    valid = [s for s in samples if len(s) > 5 and s[5] > 0]
    if not valid:
        return None, [], 0
    span = sum(s[5] for s in valid)
    first = valid[0][0] - valid[0][5]
    step = min(AVAIL_STEP, max(300, -(-int(now - first) // (48 * 60)) * 60))
    buckets = {}
    for s in valid:
        b = buckets.setdefault(int((now - s[0]) // step), [0.0, 0.0])
        b[0] += s[4]
        b[1] += s[5]
    bars = [round(1 - b[0] / b[1], 4) if b and b[1] else None
            for b in (buckets.get(i) for i in range(max(buckets), -1, -1))]
    return round(1 - sum(s[4] for s in valid) / span, 6), bars, int(now - first)

def counters(stats, server):
    c = {}
    for key in ("blocked-pro", "blocked-tif", "blocked-local", "blocked-other", "cache-hits", "cache-misses",
                "latency0-1", "latency1-10", "latency10-50", "latency50-100", "latency100-1000", "latency-slow"):
        if key in stats:
            c[key] = int(stats[key])
    for f in server.get("frontends", []):
        t = str(f.get("type", ""))
        proto = "dnscrypt" if "DNSCrypt" in t else "doh" if "HTTP" in t else "dot" if "TLS" in t else "doq" if "QUIC" in t else None
        if proto:
            c["proto-" + proto] = c.get("proto-" + proto, 0) + int(f.get("queries", 0))
    for chain in ("rules",):
        for r in server.get(chain, []):
            name = str(r.get("name", ""))
            if name in ("stats-v4", "stats-v6"):
                c[name] = int(r.get("matches", 0))
    return c

def deltas(cur, prev, restarted):
    d = {}
    for k, v in cur.items():
        if restarted or k not in prev or v < prev[k]:
            d[k] = v
        else:
            d[k] = v - prev[k]
    return d

def ops(samples, queries):
    extra = [s[6] for s in samples if len(s) > 6]
    if not extra:
        return None
    tot = {}
    for e in extra:
        for k, v in e.items():
            tot[k] = tot.get(k, 0) + v
    covered = sum(s[3] for s in samples if len(s) > 6)
    out = {"coverage": int(covered)}
    pq = {p: tot.get("proto-" + p, 0) for p in ("doh", "dot", "doq", "dnscrypt")}
    if sum(pq.values()):
        out["proto"] = {p: round(v / sum(pq.values()), 4) for p, v in pq.items()}
    if tot.get("stats-v4", 0) + tot.get("stats-v6", 0):
        out["v6"] = round(tot.get("stats-v6", 0) / (tot.get("stats-v4", 0) + tot.get("stats-v6", 0)), 4)
    if tot.get("cache-hits", 0) + tot.get("cache-misses", 0):
        out["cache_hit"] = round(tot["cache-hits"] / (tot["cache-hits"] + tot.get("cache-misses", 0)), 4)
    lb = [tot.get(k, 0) for k in ("latency0-1", "latency1-10", "latency10-50", "latency50-100", "latency100-1000", "latency-slow")]
    if sum(lb):
        out["latency_buckets"] = [round(v / sum(lb), 4) for v in lb]
    q = sum(s[1] for s in samples if len(s) > 6)
    blocked = {k: tot.get("blocked-" + k, 0) for k in ("pro", "tif", "local", "other")}
    if q and sum(blocked.values()):
        out["blocked_share"] = {k: round(v / q, 5) for k, v in blocked.items()}
        out["blocked_counts"] = blocked
    return out

def count_domains(path):
    n = 0
    with open(path, "rb") as f:
        for line in f:
            if not line or line[:1] in (b";", b"$", b"@", b" ", b"\t", b"\n", b"*"):
                continue
            n += 1
    return n

def rpz_serial(path):
    with open(path, "rb") as f:
        head = f.read(4096).decode("ascii", "replace")
    m = __import__("re").search(r"SOA\s+\S+\s+\S+\s+\(?\s*(\d+)", head)
    return int(m.group(1)) if m else None

def list_status(prev):
    rec = {}
    if REC_API:
        try:
            reckey = open(os.path.join(os.environ.get("CREDENTIALS_DIRECTORY", "/etc/dnsbunker-stats"), "recapikey")).read().strip()
            rec = get(REC_API.rstrip("/") + "/rpzstatistics", reckey)
        except Exception:
            rec = {}
    try:
        checked = int(os.stat(CHECK_STAMP).st_mtime)
    except OSError:
        checked = None
    out = {}
    for key, (path, policy) in LISTS.items():
        p = dict(prev.get(key, {}))
        try:
            st = os.stat(path)
            sig = [int(st.st_mtime), st.st_size]
            if p.get("sig") != sig:
                p.update(sig=sig, entries=count_domains(path))
                serial = rpz_serial(path)
                if serial != p.get("serial"):
                    p.update(serial=serial, updated=int(st.st_mtime))
        except OSError:
            pass
        r = rec.get(policy) if isinstance(rec, dict) else None
        if r:
            p["loaded"] = int(r.get("last_update", 0)) or p.get("loaded")
            if r.get("serial") and r["serial"] != p.get("serial"):
                p.update(serial=r["serial"], updated=int(r.get("last_update", time.time())))
        if checked:
            p["checked"] = checked
        if p:
            out[key] = p
    return out

def write(path, obj):
    tmp = path + ".tmp"
    with open(tmp, "w") as f:
        json.dump(obj, f, separators=(",", ":"))
    os.chmod(tmp, 0o644)
    os.replace(tmp, path)

stats = {s["name"]: s["value"] for s in get(API + "/statistics") if "value" in s}
server = get(API)
blocked = sum(r.get("matches", 0) for chain in ("response-rules", "cache-hit-response-rules")
              for r in server.get(chain, []) if r.get("name") in ("stats-blocked", "stats-blocked-cached"))
now, queries, uptime = time.time(), int(stats["queries"]), int(stats["uptime"])

try:
    state = json.load(open(STATE))
except (OSError, ValueError):
    state = {"samples": [], "last": None}
last, samples = state["last"], state["samples"]

backend_up = any(s.get("state") == "up" for s in server.get("servers", []))

cur = counters(stats, server)
lists = list_status(state.get("lists", {}))

if last:
    span = now - last["ts"]
    restarted = not (uptime >= last["uptime"] and queries >= last["queries"] and blocked >= last["blocked"])
    if not restarted:
        dq, db, dt = queries - last["queries"], blocked - last["blocked"], span
        down = 0.0
    else:
        dq, db, dt = queries, blocked, min(uptime, span)
        down = max(0.0, span - uptime)
    if not backend_up or dq == 0:
        down = span
    sample = [round(now), dq, db, round(dt, 3), round(down, 3), round(span, 3)]
    if "c" in last:
        sample.append(deltas(cur, last["c"], restarted))
    samples.append(sample)
    samples = [s for s in samples if s[0] > now - WINDOW]
    points, step = series(samples, now)
    avail, bars, avail_cov = availability(samples, now)
    write(OUT, {"v": 1, "updated": int(now), "queries24h": sum(s[1] for s in samples),
                "blocked24h": sum(s[2] for s in samples), "qps": round(qps_median(samples), 1),
                "coverage": int(min(WINDOW, now - samples[0][0] + 60)),
                "series": points, "series_step": step,
                "avail": avail, "avail_bars": bars, "avail_coverage": avail_cov,
                "ops": ops(samples, sum(s[1] for s in samples)),
                "lists": {k: {f: v.get(f) for f in ("entries", "updated", "loaded", "checked", "serial")} for k, v in lists.items()},
                "latency_ms": {p: round(float(stats[k]) / 1000, 2)
                               for p, k in (("doh", "latency-doh-avg1000"), ("dot", "latency-dot-avg1000"), ("doq", "latency-doq-avg1000"), ("dnscrypt", "latency-avg1000")) if k in stats}})

write(STATE, {"last": {"ts": now, "queries": queries, "blocked": blocked, "uptime": uptime, "c": cur},
              "samples": samples, "lists": lists})