#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
gpu-analyzer.py -- Pod Efficiency Analyzer, local system probe
BY: Arthur Rasmusson -- open.pod-efficiency.tools

Inspects the system it runs on (CPU, RAM, NVIDIA GPUs, NFS mounts, disk),
estimates the inference-serving uplift available with Inferra by Lightbits
Labs (KV-cache offload to fast NVMe), and uploads the statistics to
https://api.pod-efficiency.tools/api/v1/stats to enrich your PDF report.

Estimates are derived from Lightbits Labs' published LightInferra benchmarks
(4x NVIDIA L40S, ScaleFlux NVMe):
https://lightbitslabs.com/blog/introducing-lightinferra-280x-improved-ai-token-economy-by-lightbits-labs/

Usage:
    python3 gpu-analyzer.py               # analyze + upload
    python3 gpu-analyzer.py --no-upload   # analyze only, nothing leaves this box

Only Python 3 standard library is used. Nothing else is read or sent.
"""

import argparse
import json
import math
import os
import platform
import shutil
import socket
import subprocess
import sys
import urllib.request

API_URL = "https://api.pod-efficiency.tools/api/v1/stats"
BENCH_SOURCE = ("https://lightbitslabs.com/blog/introducing-lightinferra-"
                "280x-improved-ai-token-economy-by-lightbits-labs/")
FARMGPU_SOURCE = "https://blog.farmgpu.com/kv-cache-lightbits-scaleflux/"

# Published turn-2 measurements, 4x L40S rig (per-GPU figures = rig / 4).
RIG_GPUS = 4
BENCHMARKS = [
    {"model": "Qwen2.5-7B-Instruct-1M", "ctx": 100000,
     "tps_base": 5.0, "tps_inf": 95.0, "ttft_base_s": None, "ttft_inf_s": None},
    {"model": "DeepSeek-R1-Distill-Llama-70B-FP8", "ctx": 131000,
     "tps_base": 1.7, "tps_inf": 18.7, "ttft_base_s": 70.8, "ttft_inf_s": 0.465},
    {"model": "Llama-4-Scout-17B-16E-FP8", "ctx": 400000,
     "tps_base": None, "tps_inf": None, "ttft_base_s": 103.0, "ttft_inf_s": 0.457},
    {"model": "Qwen2.5-7B-Instruct-1M", "ctx": 1010000,
     "tps_base": 0.3, "tps_inf": 27.2, "ttft_base_s": 372.3, "ttft_inf_s": 1.3},
]
REUSE_SHARE = 0.70  # assumed share of multi-turn requests hitting reusable KV cache


class C:
    PLUM = "\033[35m"
    GOLD = "\033[33m"
    GREEN = "\033[32m"
    RED = "\033[31m"
    BOLD = "\033[1m"
    DIM = "\033[2m"
    END = "\033[0m"

    @classmethod
    def off(cls):
        for name in ("PLUM", "GOLD", "GREEN", "RED", "BOLD", "DIM", "END"):
            setattr(cls, name, "")


def run(cmd):
    try:
        out = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
                             timeout=20)
        return out.stdout.decode(errors="replace").strip()
    except Exception:
        return ""


def collect_cpu():
    model, cores = "unknown", os.cpu_count() or 0
    try:
        with open("/proc/cpuinfo") as f:
            for line in f:
                if line.lower().startswith("model name"):
                    model = line.split(":", 1)[1].strip()
                    break
    except OSError:
        pass
    return {"model": model, "logical_cores": cores}


def collect_mem_gib():
    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    return round(int(line.split()[1]) / 1048576.0, 1)
    except OSError:
        pass
    return None


def collect_gpus():
    if not shutil.which("nvidia-smi"):
        return []
    out = run(["nvidia-smi",
               "--query-gpu=name,memory.total,utilization.gpu,driver_version",
               "--format=csv,noheader,nounits"])
    gpus = []
    for line in out.splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) >= 4:
            gpus.append({"name": parts[0], "memory_mib": _to_num(parts[1]),
                         "utilization_pct": _to_num(parts[2]),
                         "driver": parts[3]})
    return gpus


def _to_num(s):
    try:
        return float(s)
    except ValueError:
        return None


def collect_nfs():
    mounts = []
    try:
        with open("/proc/mounts") as f:
            for line in f:
                dev, mnt, fstype = line.split()[:3]
                if fstype.startswith("nfs"):
                    mounts.append({"device": dev, "mountpoint": mnt, "type": fstype})
    except OSError:
        pass
    return mounts


def collect_disk():
    try:
        du = shutil.disk_usage("/")
        return {"total_gib": round(du.total / 2**30, 1),
                "free_gib": round(du.free / 2**30, 1)}
    except OSError:
        return None


def capacity_multiplier(k, s=REUSE_SHARE):
    return 1.0 / ((1.0 - s) + s / k) if k > 1 else 1.0


def fmt_ttft(v):
    if v is None:
        return "--"
    if v >= 60:
        return "%.1f min" % (v / 60)
    if v >= 1:
        return "%.1f s" % v
    return "%d ms" % round(v * 1000)


def print_report(info):
    p = print
    w = 78
    p("")
    p(C.PLUM + C.BOLD + "=" * w + C.END)
    p(C.PLUM + C.BOLD + "  POD EFFICIENCY ANALYZER".ljust(w - 24) + "by Arthur Rasmusson  " + C.END)
    p(C.DIM + "  Inferra by Lightbits Labs -- KV-cache offload uplift estimate" + C.END)
    p(C.PLUM + C.BOLD + "=" * w + C.END)

    p(C.BOLD + "\n  Your system" + C.END)
    p("    Host       : %s (%s %s)" % (info["hostname"], info["os"], info["kernel"]))
    p("    CPU        : %s (%s logical cores)" % (info["cpu"]["model"], info["cpu"]["logical_cores"]))
    p("    Memory     : %s GiB" % info["memory_gib"])
    if info["disk"]:
        p("    Disk (/)   : %s GiB total, %s GiB free" % (info["disk"]["total_gib"], info["disk"]["free_gib"]))
    if info["gpus"]:
        for i, g in enumerate(info["gpus"]):
            p("    GPU %-6s : %s, %.0f MiB, util %.0f%%, driver %s" % (
                "#%d" % i, g["name"], g["memory_mib"] or 0, g["utilization_pct"] or 0, g["driver"]))
    else:
        p("    GPU        : " + C.GOLD + "none detected (nvidia-smi not found or no devices)" + C.END)
    if info["nfs_mounts"]:
        for m in info["nfs_mounts"]:
            p("    NFS        : %s on %s (%s)" % (m["device"], m["mountpoint"], m["type"]))
    else:
        p("    NFS        : no NFS mounts")

    ngpus = max(len(info["gpus"]), 1)
    unit = "detected GPU(s)" if info["gpus"] else "GPU (hypothetical)"
    p(C.BOLD + "\n  Estimated multi-turn serving uplift with Inferra on your %d %s" % (ngpus, unit) + C.END)
    p(C.DIM + "  Assumes %d%% of requests reuse KV cache; M = 1/((1-s)+s/k) applied to" % (REUSE_SHARE * 100) +
      " published per-hit speedups." + C.END)
    p("")
    hdr = "    %-38s %10s %12s %12s" % ("Workload (published benchmark)", "Speedup", "tok/s now", "with Inferra")
    p(C.PLUM + hdr + C.END)
    p("    " + "-" * (w - 8))
    for b in info["uplift"]["workloads"]:
        p("    %-38s %9.0fx %12.1f %12.1f" % (
            "%s @ %dK ctx" % (b["model"].split("-Instruct")[0][:26], b["ctx"] // 1000),
            b["effective_multiplier"], b["fleet_tps_base"], b["fleet_tps_inf"]))
    p("")
    for b in info["uplift"]["ttft"]:
        p("    TTFT @ %4dK ctx : %s  ->  %s   (%dx faster)" % (
            b["ctx"] // 1000, fmt_ttft(b["ttft_base_s"]), fmt_ttft(b["ttft_inf_s"]),
            round(b["ttft_base_s"] / b["ttft_inf_s"])))
    p("    TTFT @  10M ctx  : 1,154x faster on Llama-4-Scout (FarmGPU study), plus")
    p("    3x more requests/GPU and 65% lower infrastructure cost.")

    m100k = info["uplift"]["workloads"][0]["effective_multiplier"]
    p(C.GREEN + C.BOLD + "\n  Bottom line" + C.END)
    p("    The same %d GPU(s) could serve ~%.1fx more multi-turn inference at 100K" % (ngpus, m100k))
    p("    context, or today's load could run on ~%d GPU(s) instead of %d." % (
        max(math.ceil(ngpus / m100k), 1), ngpus))
    p("    Interactive context extends to 1M+ tokens (sub-1.5 s turn-2 TTFT).")
    p(C.DIM + "\n  Sources: %s" % BENCH_SOURCE + C.END)
    p(C.DIM + "           %s" % FARMGPU_SOURCE + C.END)
    p(C.DIM + "  Estimates only -- validate with a proof of concept on your workload." + C.END)
    p(C.PLUM + C.BOLD + "=" * w + C.END)


def build_payload():
    gpus = collect_gpus()
    ngpus = max(len(gpus), 1)
    workloads, ttft = [], []
    for b in BENCHMARKS:
        if b["tps_base"]:
            k = b["tps_inf"] / b["tps_base"]
            m = capacity_multiplier(k)
            per_gpu = b["tps_base"] / RIG_GPUS
            workloads.append({
                "model": b["model"], "ctx": b["ctx"], "published_speedup": round(k, 1),
                "effective_multiplier": round(m, 1),
                "fleet_tps_base": round(per_gpu * ngpus, 1),
                "fleet_tps_inf": round(per_gpu * ngpus * m, 1),
            })
        if b["ttft_base_s"]:
            ttft.append({"ctx": b["ctx"], "ttft_base_s": b["ttft_base_s"],
                         "ttft_inf_s": b["ttft_inf_s"]})
    return {
        "analyzer_version": "0.4.0",
        "hostname": socket.gethostname(),
        "os": "%s %s" % (platform.system(), platform.release()),
        "kernel": platform.version(),
        "python": platform.python_version(),
        "cpu": collect_cpu(),
        "memory_gib": collect_mem_gib(),
        "disk": collect_disk(),
        "gpus": gpus,
        "nfs_mounts": collect_nfs(),
        "uplift": {"reuse_share": REUSE_SHARE, "workloads": workloads, "ttft": ttft,
                   "source": BENCH_SOURCE},
    }


def upload(payload):
    data = json.dumps(payload).encode()
    req = urllib.request.Request(API_URL, data=data,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=15) as resp:
        return json.loads(resp.read().decode())


def main():
    ap = argparse.ArgumentParser(description="Pod Efficiency Analyzer -- local probe")
    ap.add_argument("--no-upload", action="store_true",
                    help="analyze only; send nothing to api.pod-efficiency.tools")
    ap.add_argument("--json", action="store_true", help="print raw JSON payload too")
    args = ap.parse_args()

    if not sys.stdout.isatty():
        C.off()

    info = build_payload()
    print_report(info)
    if args.json:
        print(json.dumps(info, indent=2))

    if args.no_upload:
        print("\n  --no-upload set: nothing was sent.\n")
        return

    print("\n  Uploading statistics to %s ..." % API_URL)
    try:
        resp = upload(info)
        print(C.GREEN + "  Uploaded OK. Report id: %s" % resp.get("id") + C.END)
        print("  Reference this id on open.pod-efficiency.tools when generating your PDF.\n")
    except Exception as exc:
        print(C.RED + "  Upload failed (%s). Re-run with --no-upload to skip." % exc + C.END)
        sys.exit(1)


if __name__ == "__main__":
    main()