# plugins/host_agent/agent.py
"""Steward host agent — pushes resource metrics to a Steward instance.

Python 3.8+ stdlib only. Target ~300 lines. Served to targets at
GET /plugins/host_agent/agent.py.
"""
from __future__ import annotations

import json
import os
import re
import shutil
import signal
import socket
import sys
import time
import urllib.error
import urllib.request
from collections import deque
from datetime import datetime, timezone

AGENT_VERSION = "1.2.0"


class ConfigError(Exception):
    pass


REQUIRED_KEYS = ("url", "token")
INT_KEYS = ("interval_seconds",)
LIST_KEYS = ("mounts",)


def read_config(path: str) -> dict:
    """Parse a flat `key = value` config file."""
    cfg: dict = {}
    try:
        with open(path, "r", encoding="utf-8") as f:
            for lineno, raw in enumerate(f, 1):
                line = raw.strip()
                if not line or line.startswith("#"):
                    continue
                if "=" not in line:
                    raise ConfigError(f"{path}:{lineno}: expected 'key = value'")
                key, _, value = line.partition("=")
                key = key.strip()
                value = value.strip()
                if key in INT_KEYS:
                    try:
                        cfg[key] = int(value)
                    except ValueError:
                        raise ConfigError(f"{path}:{lineno}: {key} must be int")
                elif key in LIST_KEYS:
                    cfg[key] = [v.strip() for v in value.split(",") if v.strip()]
                else:
                    cfg[key] = value
    except FileNotFoundError:
        raise ConfigError(f"{path}: not found")

    missing = [k for k in REQUIRED_KEYS if k not in cfg]
    if missing:
        raise ConfigError(f"{path}: missing required key(s): {', '.join(missing)}")

    cfg.setdefault("interval_seconds", 30)
    return cfg


# ─── collectors ──────────────────────────────────────────────────────────────

STAT_PATH = "/proc/stat"
MEMINFO_PATH = "/proc/meminfo"
LOADAVG_PATH = "/proc/loadavg"
UPTIME_PATH = "/proc/uptime"
OS_RELEASE_PATH = "/etc/os-release"
NETDEV_PATH = "/proc/net/dev"
DISKSTATS_PATH = "/proc/diskstats"
HWMON_DIR = "/sys/class/hwmon"
PSI_DIR = "/proc/pressure"

# Partition names to drop from disk I/O (we report whole-disk throughput only).
_PARTITION_RE = re.compile(
    r"^(?:sd[a-z]+\d+|vd[a-z]+\d+|hd[a-z]+\d+|nvme\d+n\d+p\d+|mmcblk\d+p\d+)$"
)


def _read_file(path: str) -> str:
    with open(path, "r", encoding="utf-8") as f:
        return f.read()


def _parse_cpu_lines(stat_text: str) -> dict:
    """Return {name: (total_jiffies, idle_jiffies)} for every cpu line.

    'cpu' is the aggregate; 'cpu0', 'cpu1', … are per-core.
    """
    out: dict = {}
    for line in stat_text.splitlines():
        if not line.startswith("cpu"):
            continue
        parts = line.split()
        try:
            fields = [int(x) for x in parts[1:]]
        except ValueError:
            continue
        if len(fields) < 4:
            continue
        idle = fields[3] + (fields[4] if len(fields) > 4 else 0)  # idle + iowait
        out[parts[0]] = (sum(fields), idle)
    return out


def collect_cpu(sample_window: float = 0.2) -> tuple[float, list]:
    """Return (aggregate_pct, [per_core_pct, …]) over sample_window seconds."""
    a = _parse_cpu_lines(_read_file(STAT_PATH))
    time.sleep(sample_window)
    b = _parse_cpu_lines(_read_file(STAT_PATH))

    def _pct(name: str):
        if name not in a or name not in b:
            return None
        dt = b[name][0] - a[name][0]
        di = b[name][1] - a[name][1]
        if dt <= 0:
            return 0.0
        return round(100.0 * (dt - di) / dt, 2)

    agg = _pct("cpu") or 0.0
    cores: list = []
    i = 0
    while f"cpu{i}" in b:
        cores.append(_pct(f"cpu{i}"))
        i += 1
    return agg, cores


def collect_memory() -> dict:
    info: dict[str, int] = {}
    for line in _read_file(MEMINFO_PATH).splitlines():
        if ":" not in line:
            continue
        key, _, rest = line.partition(":")
        parts = rest.strip().split()
        if not parts:
            continue
        try:
            value_kb = int(parts[0])
        except ValueError:
            continue
        info[key] = value_kb * 1024  # bytes
    total = info.get("MemTotal", 0)
    available = info.get("MemAvailable", 0)
    swap_total = info.get("SwapTotal", 0)
    swap_free = info.get("SwapFree", 0)
    return {
        "total_bytes": total,
        "used_bytes": max(total - available, 0),
        "available_bytes": available,
        "swap_used_bytes": max(swap_total - swap_free, 0),
        "cached_bytes": info.get("Cached", 0),
        "buffers_bytes": info.get("Buffers", 0),
    }


def collect_storage(mounts: list[str]) -> list[dict]:
    out: list[dict] = []
    for m in mounts:
        try:
            u = shutil.disk_usage(m)
        except (FileNotFoundError, PermissionError):
            continue
        out.append({"mount": m, "total_bytes": u.total, "used_bytes": u.used})
    return out


def collect_load() -> dict:
    parts = _read_file(LOADAVG_PATH).split()
    return {"1m": float(parts[0]), "5m": float(parts[1]), "15m": float(parts[2])}


def collect_uptime() -> int:
    return int(float(_read_file(UPTIME_PATH).split()[0]))


def detect_primary_ip() -> str | None:
    """Best-effort primary non-loopback IPv4 of this host.

    connect() on a SOCK_DGRAM socket only sets the default peer and selects the
    outbound interface — no packets are sent — so getsockname() reveals the
    source IP the kernel would use to reach the internet. This is the host's own
    view of its address, so it survives reverse proxies / NAT (unlike the
    server reading request.remote_addr). Returns None on any error (offline / no
    route) and skips loopback.
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
    except OSError:
        return None
    finally:
        s.close()
    if not ip or ip.startswith("127."):
        return None
    return ip


def collect_metadata() -> dict:
    u = os.uname()
    distro = "unknown"
    try:
        for line in _read_file(OS_RELEASE_PATH).splitlines():
            if line.startswith("PRETTY_NAME="):
                distro = line.split("=", 1)[1].strip().strip('"')
                break
    except (OSError, FileNotFoundError):
        pass
    return {
        "kernel": u.release,
        "distro": distro,
        "arch": u.machine,
        "host_ip": detect_primary_ip(),
    }


def default_mounts() -> list[str]:
    """Return real mount points from /proc/mounts, skipping pseudo filesystems."""
    skip_types = {"tmpfs", "devtmpfs", "proc", "sysfs", "cgroup", "cgroup2",
                  "overlay", "squashfs", "ramfs", "devpts", "mqueue", "pstore",
                  "securityfs", "debugfs", "tracefs", "hugetlbfs", "fusectl",
                  "configfs", "bpf", "autofs", "efivarfs", "binfmt_misc"}
    mounts: list[str] = []
    try:
        with open("/proc/mounts", "r") as f:
            for line in f:
                parts = line.split()
                if len(parts) < 3:
                    continue
                if parts[2] in skip_types:
                    continue
                mounts.append(parts[1])
    except OSError:
        return ["/"]
    return mounts or ["/"]


def collect_net_raw() -> dict:
    """Return {iface: (rx_bytes, tx_bytes)} cumulative counters from /proc/net/dev.

    Skips loopback and veth* (container) interfaces to cut noise. Counters are
    monotonic; per-second rates are derived from deltas in build_sample().
    """
    out: dict = {}
    try:
        text = _read_file(NETDEV_PATH)
    except OSError:
        return out
    for line in text.splitlines():
        if ":" not in line:
            continue
        name, _, rest = line.partition(":")
        name = name.strip()
        if name == "lo" or name.startswith("veth"):
            continue
        f = rest.split()
        if len(f) < 16:
            continue
        try:
            out[name] = (int(f[0]), int(f[8]))  # rx_bytes, tx_bytes
        except ValueError:
            continue
    return out


def collect_diskio_raw() -> dict:
    """Return {device: (read_bytes, write_bytes)} cumulative from /proc/diskstats.

    Whole disks only — loop/ram/sr/fd and partitions are skipped. Sectors are
    512 bytes. Counters are monotonic; rates derived from deltas.
    """
    out: dict = {}
    try:
        text = _read_file(DISKSTATS_PATH)
    except OSError:
        return out
    for line in text.splitlines():
        f = line.split()
        if len(f) < 10:
            continue
        name = f[2]
        if name.startswith(("loop", "ram", "sr", "fd")) or _PARTITION_RE.match(name):
            continue
        try:
            out[name] = (int(f[5]) * 512, int(f[9]) * 512)  # sectors_read, sectors_written
        except ValueError:
            continue
    return out


def collect_temps() -> list:
    """Return [{label, celsius}] from /sys/class/hwmon (best-effort, may be empty)."""
    out: list = []
    try:
        chips = os.listdir(HWMON_DIR)
    except OSError:
        return out
    for hw in chips:
        d = os.path.join(HWMON_DIR, hw)
        chip = ""
        try:
            chip = _read_file(os.path.join(d, "name")).strip()
        except OSError:
            pass
        try:
            files = os.listdir(d)
        except OSError:
            continue
        for fn in files:
            if not (fn.startswith("temp") and fn.endswith("_input")):
                continue
            idx = fn[: -len("_input")]  # e.g. "temp1"
            try:
                milli = int(_read_file(os.path.join(d, fn)).strip())
            except (OSError, ValueError):
                continue
            label = ""
            try:
                label = _read_file(os.path.join(d, idx + "_label")).strip()
            except OSError:
                pass
            name = label or (f"{chip}_{idx}" if chip else idx)
            out.append({"label": name, "celsius": round(milli / 1000.0, 1)})
    return out


def _parse_psi(text: str) -> dict:
    """Parse a /proc/pressure/* file into {some_avg10, some_avg60, full_avg10, …}."""
    res: dict = {}
    for line in text.splitlines():
        parts = line.split()
        if not parts:
            continue
        kind = parts[0]  # "some" | "full"
        for p in parts[1:]:
            for prefix, suffix in (("avg10=", "avg10"), ("avg60=", "avg60")):
                if p.startswith(prefix):
                    try:
                        res[f"{kind}_{suffix}"] = float(p[len(prefix):])
                    except ValueError:
                        pass
    return res


def collect_psi() -> dict:
    """Return pressure-stall (PSI) gauges, prefixed by resource (mem/cpu/io).

    Absent on kernels without CONFIG_PSI — returns only what exists.
    """
    out: dict = {}
    for res_name, fname in (("mem", "memory"), ("cpu", "cpu"), ("io", "io")):
        try:
            text = _read_file(os.path.join(PSI_DIR, fname))
        except OSError:
            continue
        for k, v in _parse_psi(text).items():
            out[f"{res_name}_{k}"] = v
    return out


def _rates(cur: dict, prev: dict, dt: float) -> dict:
    """Per-second rates for counters present in both snapshots over dt seconds.

    Skips names missing from prev and negative deltas (counter reset / iface or
    disk hot-plug), so a reboot never emits a bogus spike.
    """
    out: dict = {}
    if dt <= 0:
        return out
    for name, cur_vals in cur.items():
        if name not in prev:
            continue
        prev_vals = prev[name]
        deltas = [c - p for c, p in zip(cur_vals, prev_vals)]
        if any(d < 0 for d in deltas):
            continue
        out[name] = tuple(d / dt for d in deltas)
    return out


# ─── ring buffer ─────────────────────────────────────────────────────────────


class RingBuffer:
    def __init__(self, maxlen: int = 20) -> None:
        self._dq: deque = deque(maxlen=maxlen)

    def push(self, item) -> None:
        self._dq.append(item)

    def drain(self) -> list:
        out = list(self._dq)
        self._dq.clear()
        return out

    def __len__(self) -> int:
        return len(self._dq)


# ─── payload ─────────────────────────────────────────────────────────────────


def build_sample(mounts: list[str], state: dict) -> dict:
    """Collect one full sample. Partial samples allowed if a collector fails.

    `state` carries the previous network/disk counters + monotonic timestamp so
    throughput rates can be derived from deltas; it is mutated in place.
    """
    sample: dict = {"ts": datetime.now(timezone.utc).isoformat()}
    try:
        agg, cores = collect_cpu()
        sample["cpu_pct"] = agg
        sample["cpu_cores"] = cores
    except Exception:
        sample["cpu_pct"] = None
    try:
        sample["mem"] = collect_memory()
    except Exception:
        sample["mem"] = None
    try:
        sample["load"] = collect_load()
    except Exception:
        sample["load"] = None
    try:
        sample["uptime_secs"] = collect_uptime()
    except Exception:
        sample["uptime_secs"] = None
    try:
        sample["storage"] = collect_storage(mounts)
    except Exception:
        sample["storage"] = []
    try:
        sample["temps"] = collect_temps()
    except Exception:
        sample["temps"] = []
    try:
        sample["psi"] = collect_psi()
    except Exception:
        sample["psi"] = {}

    # Throughput rates: diff cumulative counters against the previous sample.
    now_mono = time.monotonic()
    try:
        net_raw = collect_net_raw()
    except Exception:
        net_raw = {}
    try:
        disk_raw = collect_diskio_raw()
    except Exception:
        disk_raw = {}
    dt = now_mono - state["mono"] if state.get("mono") else 0.0
    sample["net"] = [
        {"iface": n, "rx_bps": round(rx, 1), "tx_bps": round(tx, 1)}
        for n, (rx, tx) in _rates(net_raw, state.get("net", {}), dt).items()
    ]
    sample["diskio"] = [
        {"device": n, "read_bps": round(rd, 1), "write_bps": round(wr, 1)}
        for n, (rd, wr) in _rates(disk_raw, state.get("disk", {}), dt).items()
    ]
    state["net"], state["disk"], state["mono"] = net_raw, disk_raw, now_mono
    return sample


def build_payload(samples: list[dict], hostname: str, metadata: dict) -> dict:
    return {
        "agent_version": AGENT_VERSION,
        "hostname": hostname,
        "metadata": metadata,
        "samples": samples,
    }


# ─── POST + backoff ──────────────────────────────────────────────────────────


BACKOFF_CAP = 300


def next_backoff(current: int) -> int:
    if current <= 0:
        return 30
    return min(current * 2, BACKOFF_CAP)


def post_payload(url: str, token: str, payload: dict) -> tuple[bool, int | None]:
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url.rstrip("/") + "/plugins/host_agent/ingest",
        data=body,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}",
            "User-Agent": f"steward-host-agent/{AGENT_VERSION}",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            return (200 <= resp.status < 300, resp.status)
    except urllib.error.HTTPError as e:
        return (False, e.code)
    except (urllib.error.URLError, TimeoutError, socket.timeout, OSError):
        return (False, None)


# ─── main loop ───────────────────────────────────────────────────────────────


_reload_requested = False
_shutdown_requested = False


def _handle_hup(_signum, _frame):
    global _reload_requested
    _reload_requested = True


def _handle_term(_signum, _frame):
    global _shutdown_requested
    _shutdown_requested = True


def _log(level: str, msg: str) -> None:
    sys.stderr.write(f"[{level}] {msg}\n")
    sys.stderr.flush()


def main_loop(conf_path: str) -> int:
    global _reload_requested, _shutdown_requested
    signal.signal(signal.SIGHUP, _handle_hup)
    signal.signal(signal.SIGTERM, _handle_term)

    try:
        cfg = read_config(conf_path)
    except ConfigError as e:
        _log("ERROR", str(e))
        return 2

    metadata = collect_metadata()
    hostname = cfg.get("hostname") or socket.gethostname()
    mounts = cfg.get("mounts") or default_mounts()
    buffer = RingBuffer(maxlen=20)
    backoff = 0
    # Carries previous net/disk counters + monotonic ts for rate computation.
    rate_state: dict = {}

    _log("INFO", f"steward-host-agent {AGENT_VERSION} starting "
                 f"(url={cfg['url']}, interval={cfg['interval_seconds']}s)")

    while not _shutdown_requested:
        if _reload_requested:
            try:
                cfg = read_config(conf_path)
                mounts = cfg.get("mounts") or default_mounts()
                metadata = collect_metadata()  # refresh host_ip/distro on reload
                _log("INFO", "config reloaded")
            except ConfigError as e:
                _log("ERROR", f"reload failed: {e}")
            _reload_requested = False

        sample = build_sample(mounts, rate_state)
        buffered = buffer.drain()
        payload = build_payload(
            samples=buffered + [sample],
            hostname=hostname,
            metadata=metadata,
        )
        ok, status = post_payload(cfg["url"], cfg["token"], payload)

        if ok:
            if buffered:
                _log("INFO", f"flushed {len(buffered)} buffered samples")
            backoff = 0
            sleep_for = cfg["interval_seconds"]
        else:
            if status == 400:
                _log("ERROR", "server rejected payload (400) — dropping sample")
            elif status == 401:
                _log("ERROR", "token rejected (401) — check config + UI")
                buffer.push(sample)
            else:
                _log("WARN", f"POST failed (status={status}); buffering")
                for s in buffered:
                    buffer.push(s)
                buffer.push(sample)
            backoff = next_backoff(backoff)
            sleep_for = backoff

        slept = 0.0
        while slept < sleep_for and not _shutdown_requested and not _reload_requested:
            time.sleep(min(1.0, sleep_for - slept))
            slept += 1.0

    _log("INFO", "SIGTERM — flushing and exiting")
    final = buffer.drain()
    if final:
        post_payload(cfg["url"], cfg["token"],
                     build_payload(final, hostname, metadata))
    return 0


if __name__ == "__main__":
    conf = os.environ.get("STEWARD_AGENT_CONFIG", "/etc/steward-agent.conf")
    sys.exit(main_loop(conf))