#!/usr/bin/env python3
"""Restricted network-profile helper for the edge host."""
from __future__ import annotations
import json, os, re, shlex, sqlite3, subprocess, sys, time

DB = "/var/lib/streamscope/control.db"
DEV = "ens3"
CLASSES = {"ch": "1:200", "vpn": "1:300", "default": "1:999"}
SERVICES = ("netqos", "streamscope-worker", "streamscope-admin", "promtail")

def run(argv, timeout=30):
    try:
        p = subprocess.run(argv, capture_output=True, text=True, timeout=timeout)
        return p.returncode, p.stdout.strip(), p.stderr.strip()
    except Exception as exc:
        return 125, "", str(exc)

def db():
    c = sqlite3.connect(DB, timeout=15)
    c.row_factory = sqlite3.Row
    c.execute("PRAGMA busy_timeout=15000")
    return c

def rate_map():
    rc, out, _ = run(["sudo", "tc", "class", "show", "dev", DEV])
    found = {}
    for line in out.splitlines():
        m = re.search(r"^class htb (1:\w+).*?\brate\s+([0-9.]+)([KMG])bit", line, re.I)
        if not m:
            continue
        value = float(m.group(2)) * {"K": .001, "M": 1, "G": 1000}[m.group(3).upper()]
        target = next((k for k, v in CLASSES.items() if v == m.group(1)), m.group(1))
        found[target] = value
    return found, rc


def tc_metrics():
    rc, out, _ = run(["sudo", "tc", "-s", "class", "show", "dev", DEV])
    result = {}
    current = None
    for line in out.splitlines():
        m = re.search(r"^class htb (1:\w+).*?\brate\s+([0-9.]+)([KMG])bit", line, re.I)
        if m:
            current = next((key for key, value in CLASSES.items() if value == m.group(1)), m.group(1))
            rate = float(m.group(2)) * {"K": .001, "M": 1, "G": 1000}[m.group(3).upper()]
            result[current] = {"rate_mbit": rate, "bytes": 0, "packets": 0, "drops": 0}
            continue
        if current:
            sent = re.search(r"Sent\s+(\d+)\s+bytes\s+(\d+)\s+pkt.*?dropped\s+(\d+)", line)
            if sent:
                result[current].update(bytes=int(sent.group(1)), packets=int(sent.group(2)), drops=int(sent.group(3)))
                current = None
    return result, rc


def host_metrics():
    up = int(float(open("/proc/uptime").read().split()[0]))
    load = [round(value, 2) for value in os.getloadavg()]
    mem = {}
    for line in open("/proc/meminfo"):
        key, value = line.split(":", 1)
        mem[key] = int(value.strip().split()[0])
    disk = os.statvfs("/")
    return {
        "uptime_seconds": up,
        "load": load,
        "memory_total": mem.get("MemTotal", 0) * 1024,
        "memory_used": (mem.get("MemTotal", 0) - mem.get("MemAvailable", 0)) * 1024,
        "disk_total": disk.f_blocks * disk.f_frsize,
        "disk_used": (disk.f_blocks - disk.f_bavail) * disk.f_frsize,
    }

def state():
    c = db()
    limits = {r["target"]: r["limit_mbit"] for r in c.execute("SELECT target,limit_mbit FROM global_limits")}
    settings = {r["key"]: r["value"] for r in c.execute("SELECT key,value FROM settings")}
    active_sessions = c.execute("SELECT count(*) FROM session_history WHERE active=1").fetchone()[0]
    session_rows = [dict(r) for r in c.execute("SELECT user_name,device,client,ip,title,series,started_at,last_seen,country,city,asn,bytes FROM session_history WHERE active=1 ORDER BY last_seen DESC LIMIT 20")]
    recent = [dict(r) for r in c.execute("SELECT ts,user_name,kind,detail FROM multisession_events ORDER BY id DESC LIMIT 8")]
    active_blocks = c.execute("SELECT count(*) FROM multisession_blocks WHERE active=1 AND expires_at>strftime('%s','now')").fetchone()[0]
    totals = {
        "sessions": c.execute("SELECT count(*) FROM session_history").fetchone()[0],
        "known_ips": c.execute("SELECT count(*) FROM ip_history").fetchone()[0],
        "policy_events": c.execute("SELECT count(*) FROM multisession_events").fetchone()[0],
        "blocks": c.execute("SELECT count(*) FROM multisession_blocks").fetchone()[0],
    }
    c.close()
    tc_detail, tc_rc = tc_metrics()
    tc = {key: value["rate_mbit"] for key, value in tc_detail.items()}
    ipsets = {}
    for name in ("geo-a-v4", "geo-a-v6", "cloud-v4", "cloud-v6", "persist-v4", "persist-v6"):
        rc, out, _ = run(["sudo", "ipset", "list", name, "-t"])
        m = re.search(r"Number of entries:\s*(\d+)", out)
        ipsets[name] = int(m.group(1)) if m else 0
    services = {}
    for name in SERVICES:
        rc, out, _ = run(["systemctl", "is-active", name])
        services[name] = out or "unknown"
    rc, rules, _ = run(["sudo", "iptables", "-t", "mangle", "-S"])
    return {
        "ok": tc_rc == 0 and all(v == "active" for k, v in services.items() if k != "netqos"),
        "ts": int(time.time()), "host": host_metrics(), "limits": limits, "tc": tc, "tc_detail": tc_detail, "ipsets": ipsets,
        "services": services, "active_sessions": active_sessions, "sessions": session_rows,
        "active_blocks": active_blocks, "totals": totals,
        "policy": {
            "enabled": settings.get("multisession_enabled", "0") == "1",
            "mode": settings.get("multisession_mode", "monitor"),
            "min_sessions": int(settings.get("multisession_min_sessions", "2")),
            "block_minutes": int(settings.get("multisession_block_minutes", "60")),
            "persist_seconds": int(settings.get("multisession_persist_seconds", "60")),
        },
        "mark_rules": sum(1 for line in rules.splitlines() if "--set-xmark" in line or "--set-mark" in line),
        "recent_events": recent,
    }


def history():
    c = db()
    data = {
        "sessions": [dict(r) for r in c.execute(
            "SELECT user_name,device,client,ip,title,series,started_at,ended_at,last_seen,bytes,country,city,asn,active "
            "FROM session_history ORDER BY last_seen DESC LIMIT 150"
        )],
        "ips": [dict(r) for r in c.execute(
            "SELECT user_name,ip,device,client,first_seen,last_seen,sightings,limited_sightings,last_classification,country,city,asn "
            "FROM ip_history ORDER BY last_seen DESC LIMIT 150"
        )],
        "events": [dict(r) for r in c.execute(
            "SELECT ts,user_name,kind,detail FROM multisession_events ORDER BY id DESC LIMIT 150"
        )],
        "blocks": [dict(r) for r in c.execute(
            "SELECT ip,user_name,reason,blocked_at,expires_at,active,manual_allow_until FROM multisession_blocks ORDER BY blocked_at DESC LIMIT 150"
        )],
    }
    c.close()
    return {"ok": True, "history": data, "generated_at": int(time.time())}


def set_limits(ch, vpn):
    values = {"ch": float(ch), "vpn": float(vpn)}
    if any(not 0.1 <= value <= 100 for value in values.values()):
        raise ValueError("rate outside 0.1..100 Mbit")
    old, _ = rate_map()
    changed = []
    try:
        for target, value in values.items():
            cls = CLASSES[target]
            rc, out, err = run(["sudo", "tc", "class", "change", "dev", DEV, "parent", "1:", "classid", cls, "htb", "rate", f"{value:g}mbit", "ceil", f"{value:g}mbit"])
            if rc:
                raise RuntimeError(err or out or f"tc failed for {target}")
            changed.append(target)
        c = db()
        for target, value in values.items():
            c.execute("UPDATE global_limits SET limit_mbit=?,updated_at=datetime('now') WHERE target=?", (value, target))
        c.execute("UPDATE limit_profiles SET active=0")
        c.commit(); c.close()
    except Exception:
        for target in changed:
            if target in old:
                value = old[target]; cls = CLASSES[target]
                run(["sudo", "tc", "class", "change", "dev", DEV, "parent", "1:", "classid", cls, "htb", "rate", f"{value:g}mbit", "ceil", f"{value:g}mbit"])
        raise
    return {"ok": True, "limits": values, "tc": rate_map()[0]}

def set_policy(mode, min_sessions=None, block_minutes=None):
    if mode not in {"off", "monitor", "block"}:
        raise ValueError("invalid policy mode")
    c = db()
    c.execute("UPDATE settings SET value=? WHERE key='multisession_enabled'", ("0" if mode == "off" else "1",))
    if mode != "off":
        c.execute("UPDATE settings SET value=? WHERE key='multisession_mode'", (mode,))
    if min_sessions is not None:
        value = int(min_sessions)
        if not 2 <= value <= 8: raise ValueError("min sessions outside 2..8")
        c.execute("UPDATE settings SET value=? WHERE key='multisession_min_sessions'", (str(value),))
    if block_minutes is not None:
        value = int(block_minutes)
        if not 1 <= value <= 1440: raise ValueError("block minutes outside 1..1440")
        c.execute("UPDATE settings SET value=? WHERE key='multisession_block_minutes'", (str(value),))
    c.commit(); c.close()
    rc, out, err = run(["sudo", "systemctl", "restart", "streamscope-worker"], timeout=20)
    if rc: raise RuntimeError(err or out or "worker restart failed")
    return {"ok": True, "policy": state()["policy"]}

def background_limiter(action):
    if action not in {"start", "stop", "refresh"}: raise ValueError("invalid limiter action")
    unit = "netprofile-rebuild" if action in {"start", "refresh"} else "netprofile-clear"
    rc, out, err = run(["sudo", "systemd-run", "--unit", unit, "--collect", "/usr/lib/systemd/system/netqos/link-optimizer.sh", action], timeout=15)
    if rc and "already exists" not in (out + err): raise RuntimeError(err or out)
    return {"ok": True, "started": action}

def restart(name):
    mapping = {"worker": "streamscope-worker", "admin": "streamscope-admin"}
    if name not in mapping: raise ValueError("invalid service")
    rc, out, err = run(["sudo", "systemctl", "restart", mapping[name]], timeout=20)
    if rc: raise RuntimeError(err or out)
    return {"ok": True, "restarted": name}

def main():
    raw = os.environ.get("SSH_ORIGINAL_COMMAND", "state").strip()
    argv = shlex.split(raw)
    if not argv: argv = ["state"]
    if argv == ["state"]:
        result = state()
    elif argv == ["history"]:
        result = history()
    elif len(argv) == 3 and argv[0] == "set-limits":
        result = set_limits(argv[1], argv[2])
    elif argv[0] == "set-policy" and 2 <= len(argv) <= 4:
        result = set_policy(argv[1], argv[2] if len(argv) > 2 else None, argv[3] if len(argv) > 3 else None)
    elif len(argv) == 2 and argv[0] == "limiter":
        result = background_limiter(argv[1])
    elif len(argv) == 2 and argv[0] == "restart":
        result = restart(argv[1])
    else:
        raise ValueError("unsupported command")
    print(json.dumps(result, ensure_ascii=False, separators=(",", ":")))

if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(json.dumps({"ok": False, "error": str(exc)[:300]}, ensure_ascii=False, separators=(",", ":")))
        sys.exit(1)