#!/usr/bin/env python3
import json, subprocess, time, os, sys, math, shutil
from datetime import datetime, timezone, timedelta

STATE_DIR = "/var/lib/snowpanel"
STATS = os.path.join(STATE_DIR, "stats.json")
META  = os.path.join(STATE_DIR, "meta.json")
CFG   = "/etc/snowpanel/app.json"

def sh(cmd):
    return subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True).stdout.strip()

def load_json(path, default):
    try:
        with open(path, "r") as f: return json.load(f)
    except Exception:
        return default

def save_json(path, obj):
    tmp = path + ".tmp"
    with open(tmp, "w") as f:
        json.dump(obj, f, separators=(",", ":"), ensure_ascii=False)
    os.replace(tmp, path)

def now():
    return int(time.time())

def service_bytes():
    out = sh(["/bin/systemctl","show","-p","IPIngressBytes","-p","IPEgressBytes","snowflake-proxy"])
    rx = tx = 0
    for line in out.splitlines():
        if line.startswith("IPIngressBytes="):
            try: rx = int(line.split("=",1)[1] or "0")
            except: pass
        elif line.startswith("IPEgressBytes="):
            try: tx = int(line.split("=",1)[1] or "0")
            except: pass
    return rx, tx

def period_start_for_reset_day(reset_day: int) -> int:
    reset_day = max(1, min(28, int(reset_day or 1)))
    now_dt = datetime.now(timezone.utc).astimezone()
    year = now_dt.year
    month = now_dt.month
    this_start = datetime(year, month, reset_day, 0, 0, 0, tzinfo=now_dt.tzinfo)
    if now_dt >= this_start:
        start = this_start
    else:
        if month == 1:
            year -= 1; month = 12
        else:
            month -= 1
        start = datetime(year, month, reset_day, 0, 0, 0, tzinfo=now_dt.tzinfo)
    return int(start.timestamp())

def main():
    os.makedirs(STATE_DIR, exist_ok=True)

    rx, tx = service_bytes()
    t = now()

    stats = load_json(STATS, {"data":[]})
    arr = stats.get("data", [])
    arr.append({"t": t, "read": int(rx), "written": int(tx)})
    if len(arr) > 5000:
        arr = arr[-5000:]
    stats["data"] = arr
    save_json(STATS, stats)

    cfg = load_json(CFG, {})
    cap_gb        = int(cfg.get("cap_gb", 0))
    cap_reset_day = int(cfg.get("cap_reset_day", 1))
    rate_mbps     = int(cfg.get("rate_mbps", 0))

    start_ts = period_start_for_reset_day(cap_reset_day)

    rx_sum = 0
    tx_sum = 0
    prev = None
    for point in arr:
        if point["t"] < start_ts:
            continue
        if prev is not None:
            dt = point["t"] - prev["t"]
            if dt <= 0 or dt > 3600:
                prev = point; continue
            dr = max(0, point["read"]    - prev["read"])
            dw = max(0, point["written"] - prev["written"])
            rx_sum += dr
            tx_sum += dw
        prev = point

    total = rx_sum + tx_sum
    cap_bytes = cap_gb * 1024 * 1024 * 1024 if cap_gb > 0 else 0

    current_rate_mbps = 0.0
    if len(arr) >= 2:
        a = arr[-2]; b = arr[-1]
        dt = max(1, b["t"] - a["t"])
        dr = max(0, b["read"]    - a["read"])
        dw = max(0, b["written"] - a["written"])
        current_rate_mbps = ((dr + dw) * 8.0) / dt / 1_000_000.0

    label = datetime.fromtimestamp(start_ts).strftime("%Y-%m-%d") + f" (reset day {cap_reset_day})"
    meta = {
        "start_ts": start_ts,
        "period_label": label,
        "rx": rx_sum,
        "tx": tx_sum,
        "total": total,
        "cap_bytes": cap_bytes,
        "cap_hit": False,
        "current_rate_mbps": current_rate_mbps,
        "rate_mbps": rate_mbps
    }

    if cap_bytes and total >= cap_bytes:
        meta["cap_hit"] = True
        active = sh(["/bin/systemctl","is-active","snowflake-proxy"]) == "active"
        if active:
            subprocess.run(["/usr/bin/sudo","/bin/systemctl","stop","snowflake-proxy"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    save_json(META, meta)

if __name__ == "__main__":
    main()