# Phase 0 CPU probe B. Follows up the one gap in p0_cpu_probe.py: its /proc/meminfo parser
# crashed (E-003), so instance disk and RAM -- the numbers §3.13 depends on -- were never read.
# Also answers a question the first probe raised: the CPU image ships torch 2.10.0+cpu, so
# which image/torch a GPU session actually gets has to be observed from inside one.
# Zero GPU quota. No credentials.

import json
import os
import shutil
import subprocess
import sys
import time

t0 = time.time()


def meminfo():
    out = {}
    with open("/proc/meminfo") as f:
        for line in f:
            if ":" not in line:
                continue
            k, rest = line.split(":", 1)
            parts = rest.split()
            if parts:
                out[k] = int(parts[0]) // 1024
    return {k: out[k] for k in ("MemTotal", "MemAvailable", "MemFree", "SwapTotal", "Shmem") if k in out}


def disk():
    out = {}
    for p in ["/kaggle/working", "/kaggle/input", "/tmp", "/.cache", "/root/.cache", "/dev/shm", "/usr"]:
        try:
            u = shutil.disk_usage(p)
            out[p] = {"total_GB": round(u.total / 1024**3, 1), "used_GB": round(u.used / 1024**3, 1),
                      "free_GB": round(u.free / 1024**3, 1)}
        except Exception as e:
            out[p] = f"missing ({type(e).__name__})"
    return out


def cgroup():
    out = {}
    for name, paths in {
        "mem_max": ["/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"],
        "cpu_max": ["/sys/fs/cgroup/cpu.max", "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"],
    }.items():
        for p in paths:
            if os.path.exists(p):
                out[name] = open(p).read().strip()
                out[name + "_path"] = p
                break
        else:
            out[name] = "not found"
    return out


R = {
    "shape": {"CUDA_VISIBLE_DEVICES_present": "CUDA_VISIBLE_DEVICES" in os.environ,
              "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES"),
              "KAGGLE_DOCKER_IMAGE": os.environ.get("KAGGLE_DOCKER_IMAGE"),
              "KAGGLE_KERNEL_RUN_TYPE": os.environ.get("KAGGLE_KERNEL_RUN_TYPE"),
              "KAGGLE_CONTAINER_NAME": os.environ.get("KAGGLE_CONTAINER_NAME"),
              "KAGGLE_GCP_ZONE": os.environ.get("KAGGLE_GCP_ZONE"),
              "KAGGLE_KERNEL_INTEGRATIONS": os.environ.get("KAGGLE_KERNEL_INTEGRATIONS")},
    "cpu": {"os_cpu_count": os.cpu_count(), "affinity": len(os.sched_getaffinity(0)),
            "loadavg": os.getloadavg(),
            "model": [l for l in open("/proc/cpuinfo") if "model name" in l][:1]},
    "mem_MB": meminfo(),
    "cgroup": cgroup(),
    "disk": disk(),
    "df_h": subprocess.run("df -h", shell=True, capture_output=True, text=True).stdout,
    "du_cache": subprocess.run("du -sh /root/.cache/* 2>/dev/null | head -10", shell=True,
                               capture_output=True, text=True).stdout,
}
import socket

dns = {}
for h in ["cdn-lfs.huggingface.co", "cdn-lfs-us-1.huggingface.co", "huggingface.co", "xethub.huggingface.co"]:
    try:
        dns[h] = sorted({i[4][0] for i in socket.getaddrinfo(h, 443)})[:2]
    except Exception as e:
        dns[h] = f"DNSFAIL {type(e).__name__}"
R["dns"] = dns
R["_seconds"] = round(time.time() - t0, 1)
print("PROBE_JSON_BEGIN")
print(json.dumps(R, indent=1, default=str))
print("PROBE_JSON_END")