#!/usr/bin/env python3
"""Merge container resource snapshot into landing stats.json performance block.

Usage:
  merge_resources.py STATS.json RESOURCES.json [-o OUT.json]

RESOURCES shape (from mem_snapshot_emit / collect_container_resources.sh):
  { "ok": true, "loadavg": "0.1,0.2,0.3", "memory": { ... } }
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict


def compact_memory(mem: Dict[str, Any]) -> Dict[str, Any]:
    """Ensure human labels + short top cmd for UI width."""
    if not isinstance(mem, dict):
        return mem
    # Prefer explicit label (with limit); else human
    if not mem.get("container_rss_label"):
        h = mem.get("container_rss_human") or "—"
        lim = mem.get("cgroup_limit_human")
        if lim:
            mem["container_rss_label"] = f"{h} / {lim}"
        else:
            mem["container_rss_label"] = h
    tops = mem.get("top")
    if isinstance(tops, list):
        for t in tops:
            if not isinstance(t, dict):
                continue
            cmd = str(t.get("cmd") or "")
            if len(cmd) > 72:
                t["cmd_full"] = cmd
                t["cmd"] = cmd[:69] + "…"
            # ensure human present
            if not t.get("rss_human") and t.get("rss_bytes") is not None:
                try:
                    b = int(t["rss_bytes"])
                    if b < 1024:
                        t["rss_human"] = f"{b} B"
                    elif b < 1048576:
                        t["rss_human"] = f"{b/1024:.1f} KiB"
                    elif b < 1073741824:
                        t["rss_human"] = f"{b/1048576:.1f} MiB"
                    else:
                        t["rss_human"] = f"{b/1073741824:.2f} GiB"
                except Exception:
                    pass
    return mem


def merge(stats: Dict[str, Any], resources: Dict[str, Any]) -> Dict[str, Any]:
    perf = stats.setdefault("performance", {})
    if not isinstance(perf, dict):
        perf = {}
        stats["performance"] = perf

    if resources.get("loadavg"):
        perf["loadavg"] = resources["loadavg"]
        # rootless: /proc/loadavg is almost always the host's
        perf["loadavg_source"] = resources.get("loadavg_source") or "host_via_container_proc"
        if resources.get("loadavg_note"):
            perf["loadavg_note"] = resources["loadavg_note"]

    if resources.get("cpu_percent") is not None:
        perf["cpu_percent"] = resources["cpu_percent"]
        perf["cpu_percent_label"] = resources.get("cpu_percent_label") or (
            f'{resources["cpu_percent"]:.2f}%'
        )

    if resources.get("pids") is not None:
        perf["pids"] = resources["pids"]

    mem = resources.get("memory")
    if isinstance(mem, dict) and mem:
        m = compact_memory(dict(mem))
        # normalize empty groups for UI
        for key in ("taler", "java", "nginx", "postgres", "redis", "other"):
            try:
                n = int(m.get(f"{key}_n") or 0)
                b = int(m.get(f"{key}_rss_bytes") or 0)
            except Exception:
                n, b = 0, 0
            if n <= 0 or b <= 0:
                m[f"{key}_n"] = 0
                m[f"{key}_rss_bytes"] = 0
                m[f"{key}_rss_human"] = None
        h = m.get("container_rss_human") or "—"
        has_cgroup = bool(m.get("cgroup_bytes"))
        if m.get("cgroup_limit_human") and h != "—":
            m["container_rss_label"] = f"{h} / {m['cgroup_limit_human']}"
            m["memory_basis"] = "cgroup"
        elif not has_cgroup and h != "—":
            m["container_rss_label"] = f"{h} (proc)"
            m["memory_basis"] = m.get("memory_basis") or "proc_rss_sum"
        elif not m.get("container_rss_label"):
            m["container_rss_label"] = h
        m["source"] = resources.get("source") or "mem-snapshot"
        perf["memory"] = m
    return stats


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("stats")
    ap.add_argument("resources")
    ap.add_argument("-o", "--out", default="")
    args = ap.parse_args()

    stats = json.loads(Path(args.stats).read_text(encoding="utf-8"))
    resources = json.loads(Path(args.resources).read_text(encoding="utf-8"))
    if not isinstance(stats, dict) or not stats.get("ok"):
        print("skip: stats not ok", file=sys.stderr)
        return 1
    if not isinstance(resources, dict) or resources.get("ok") is False:
        print("skip: resources not ok", file=sys.stderr)
        return 1

    merge(stats, resources)
    out = Path(args.out) if args.out else Path(args.stats)
    tmp = out.with_suffix(out.suffix + f".tmp.{os.getpid()}")
    tmp.write_text(json.dumps(stats, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    tmp.replace(out)
    mem = (stats.get("performance") or {}).get("memory") or {}
    print(
        f"merged memory container={mem.get('container_rss_human')} "
        f"pg={mem.get('postgres_rss_human')} loadavg={stats.get('performance', {}).get('loadavg')}",
        file=sys.stderr,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())