# plugins/host_agent/agent.py
"""Roundtable host agent — pushes resource metrics to a Roundtable 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 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.0.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"


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


def _parse_cpu_line(stat_text: str) -> tuple[int, int]:
    """Return (total_jiffies, idle_jiffies) for the aggregate cpu line."""
    for line in stat_text.splitlines():
        if line.startswith("cpu "):
            parts = line.split()
            fields = [int(x) for x in parts[1:]]
            idle = fields[3] + (fields[4] if len(fields) > 4 else 0)  # idle + iowait
            total = sum(fields)
            return total, idle
    raise RuntimeError("no aggregate cpu line in /proc/stat")


def collect_cpu(sample_window: float = 0.2) -> float:
    """Return CPU utilization % over sample_window seconds."""
    t1, i1 = _parse_cpu_line(_read_file(STAT_PATH))
    time.sleep(sample_window)
    t2, i2 = _parse_cpu_line(_read_file(STAT_PATH))
    dt = t2 - t1
    di = i2 - i1
    if dt <= 0:
        return 0.0
    return round(100.0 * (dt - di) / dt, 2)


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),
    }


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 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}


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 ["/"]


# ─── 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]) -> dict:
    """Collect one full sample. Partial samples allowed if a collector fails."""
    sample: dict = {"ts": datetime.now(timezone.utc).isoformat()}
    try:
        sample["cpu_pct"] = collect_cpu()
    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"] = []
    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"roundtable-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

    _log("INFO", f"roundtable-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()
                _log("INFO", "config reloaded")
            except ConfigError as e:
                _log("ERROR", f"reload failed: {e}")
            _reload_requested = False

        sample = build_sample(mounts)
        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("ROUNDTABLE_AGENT_CONFIG", "/etc/roundtable-agent.conf")
    sys.exit(main_loop(conf))