#!/usr/bin/env python3
"""
mycel-node — readiness check for the MYCEL substrate.

Reads what is actually on this machine, measures what can be measured
without pulling in a toolchain, and tells you what the card would earn
at the substrate's current published rate.

It does not phone home. The substrate is not accepting nodes yet, so
there is nothing to connect to — this is the check you run before there is.

    python mycel_node.py            full report
    python mycel_node.py --json     machine-readable
    python mycel_node.py --bench    also run a FLOPS benchmark (needs torch)

No dependencies beyond the standard library.
"""

import argparse
import hashlib
import json
import platform
import shutil
import subprocess
import sys
import time

VERSION = "0.3.1"

# Rates the substrate publishes, in MYCEL per GPU-hour. Cards not listed
# fall back to the nearest tier by VRAM.
RATES = {
    "H200": 12050, "H100": 8400, "MI300X": 9700, "A100": 4600,
    "L40S": 3050, "L40": 2700, "A6000": 2100, "5090": 1700,
    "4090": 1050, "4080": 750, "3090": 600, "A5000": 700,
    "3080": 450, "3070": 310, "3060": 240, "4070": 525, "4060": 320,
    "2080": 270, "V100": 1300, "T4": 350,
}
TIERS = [(140, 10000), (78, 4500), (46, 2900), (30, 1550), (22, 950), (14, 450), (0, 210)]

RESET, DIM, GOLD, CYAN, WHITE, RED = (
    "\033[0m", "\033[2m", "\033[38;5;180m", "\033[38;5;51m", "\033[97m", "\033[38;5;203m"
)


def _unicode_ok() -> bool:
    """Windows consoles still default to cp1252. Ask for UTF-8, and fall back
    to ASCII box drawing rather than crashing on a dash."""
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        return True
    except Exception:
        pass
    return "utf" in (getattr(sys.stdout, "encoding", "") or "").lower()


UNI = _unicode_ok()
RULE = "─" if UNI else "-"
FULL = "█" if UNI else "#"
EMPTY = "░" if UNI else "."
DOT = "·" if UNI else "-"


def supports_colour() -> bool:
    if not sys.stdout.isatty():
        return False
    if platform.system() == "Windows":
        try:
            import ctypes
            k = ctypes.windll.kernel32
            k.SetConsoleMode(k.GetStdHandle(-11), 7)
            return True
        except Exception:
            return False
    return True


C = supports_colour()
def c(text, colour):
    return f"{colour}{text}{RESET}" if C else text


def run(cmd):
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=20)
        return out.stdout.strip() if out.returncode == 0 else None
    except Exception:
        return None


def nvidia_gpus():
    if not shutil.which("nvidia-smi"):
        return []
    fields = ("name,memory.total,driver_version,uuid,power.limit,"
              "clocks.max.graphics,pcie.link.width.max,temperature.gpu,utilization.gpu")
    raw = run(["nvidia-smi", f"--query-gpu={fields}", "--format=csv,noheader,nounits"])
    if not raw:
        return []
    gpus = []
    for line in raw.splitlines():
        p = [x.strip() for x in line.split(",")]
        if len(p) < 9:
            continue
        gpus.append({
            "vendor": "NVIDIA", "name": p[0], "vram_mb": to_int(p[1]),
            "driver": p[2], "uuid": p[3], "power_w": to_int(p[4]),
            "clock_mhz": to_int(p[5]), "pcie_width": to_int(p[6]),
            "temp_c": to_int(p[7]), "util_pct": to_int(p[8]),
        })
    return gpus


def amd_gpus():
    if not shutil.which("rocm-smi"):
        return []
    raw = run(["rocm-smi", "--showproductname", "--showmeminfo", "vram", "--json"])
    if not raw:
        return []
    try:
        data = json.loads(raw)
    except Exception:
        return []
    gpus = []
    for key, val in data.items():
        name = val.get("Card series") or val.get("Card model") or "AMD GPU"
        vram = val.get("VRAM Total Memory (B)")
        gpus.append({
            "vendor": "AMD", "name": str(name).strip(),
            "vram_mb": int(vram) // (1024 * 1024) if vram else 0,
            "driver": val.get("Driver version", "?"), "uuid": key,
            "power_w": 0, "clock_mhz": 0, "pcie_width": 0,
            "temp_c": 0, "util_pct": 0,
        })
    return gpus


def to_int(v):
    try:
        return int(float(v))
    except Exception:
        return 0


def rate_for(name: str, vram_mb: int) -> tuple:
    """Returns (rate, how it was decided)."""
    upper = name.upper().replace(" ", "")
    for key, rate in RATES.items():
        if key in upper:
            return rate, "listed"
    gb = vram_mb / 1024
    for floor, rate in TIERS:
        if gb >= floor:
            return rate, "tier"
    return 420, "tier"


def node_id(gpus) -> str:
    """Deterministic from the hardware itself, so the same box always
    produces the same id and two boxes never collide."""
    seed = "|".join(sorted(g["uuid"] for g in gpus)) or platform.node()
    return "MY-" + hashlib.sha256(seed.encode()).hexdigest()[:10].upper()


def benchmark(gpus):
    """Real FP16 matmul throughput. Skipped unless torch is present —
    we would rather report nothing than a number we did not measure."""
    try:
        import torch
    except ImportError:
        return None
    if not torch.cuda.is_available():
        return None
    results = []
    for i in range(min(len(gpus), torch.cuda.device_count())):
        dev = torch.device(f"cuda:{i}")
        n = 4096
        a = torch.randn(n, n, device=dev, dtype=torch.float16)
        b = torch.randn(n, n, device=dev, dtype=torch.float16)
        for _ in range(3):                       # warm up
            a @ b
        torch.cuda.synchronize()
        t = time.perf_counter()
        loops = 12
        for _ in range(loops):
            a @ b
        torch.cuda.synchronize()
        secs = time.perf_counter() - t
        tflops = (2 * n ** 3 * loops) / secs / 1e12
        results.append(round(tflops, 1))
        del a, b
        torch.cuda.empty_cache()
    return results


def collect(bench=False):
    gpus = nvidia_gpus() + amd_gpus()
    flops = benchmark(gpus) if bench and gpus else None
    for i, g in enumerate(gpus):
        g["rate"], g["rate_basis"] = rate_for(g["name"], g["vram_mb"])
        if flops and i < len(flops):
            g["tflops_fp16"] = flops[i]
    return {
        "client": f"mycel-node/{VERSION}",
        "node_id": node_id(gpus) if gpus else None,
        "host": {"os": platform.system(), "release": platform.release(),
                 "python": platform.python_version(), "cpu": platform.processor() or "?"},
        "gpus": gpus,
        "hourly_mycel": sum(g["rate"] for g in gpus),
        "substrate": "not accepting nodes yet",
    }


def bar(pct, width=22):
    filled = round(max(0, min(100, pct)) / 100 * width)
    return FULL * filled + EMPTY * (width - filled)


def report(data):
    line = RULE * 58
    print()
    print(c("  MYCEL", GOLD) + c("  node readiness check", DIM))
    print(c("  " + line, DIM))

    gpus = data["gpus"]
    if not gpus:
        print()
        print("  " + c("No supported GPU found.", RED))
        print(c("  Needs nvidia-smi (NVIDIA) or rocm-smi (AMD) on PATH.", DIM))
        print(c("  A card the driver cannot see is a card the substrate cannot rent.", DIM))
        print()
        return 1

    print()
    print(f"  {c('NODE', DIM)}   {c(data['node_id'], WHITE)}")
    print(f"  {c('HOST', DIM)}   {data['host']['os']} {data['host']['release']}"
          f"  {DOT}  python {data['host']['python']}")
    print()

    for i, g in enumerate(gpus):
        print(f"  {c(f'[{i}]', DIM)} {c(g['name'], WHITE)}")
        print(f"      {c('VRAM', DIM)}     {g['vram_mb'] / 1024:.0f} GB"
              f"        {c('DRIVER', DIM)}  {g['driver']}")
        if g["clock_mhz"]:
            print(f"      {c('CLOCK', DIM)}    {g['clock_mhz']} MHz"
                  f"     {c('POWER', DIM)}   {g['power_w']} W"
                  f"     {c('PCIE', DIM)}  x{g['pcie_width']}")
        if g["temp_c"]:
            print(f"      {c('TEMP', DIM)}     {g['temp_c']}°C"
                  f"         {c('LOAD', DIM)}    {bar(g['util_pct'])} {g['util_pct']}%")
        if "tflops_fp16" in g:
            print(f"      {c('MEASURED', DIM)} {c(str(g['tflops_fp16']) + ' TFLOPS', CYAN)}"
                  f" {c('fp16 dense matmul', DIM)}")
        basis = "published rate for this card" if g["rate_basis"] == "listed" else "tier estimate from VRAM"
        rate_txt = f"{g['rate']:,} MYCEL/hr"
        print(f"      {c('RATE', DIM)}     {c(rate_txt, GOLD)}  {c(DOT + ' ' + basis, DIM)}")
        print()

    print(c("  " + line, DIM))
    hourly = data["hourly_mycel"]
    print(f"  {c('AT FULL LOAD', DIM)}  {c(f'{hourly:,}', GOLD)} {c('MYCEL/hr', DIM)}"
          f"   {DOT}   {c(f'{hourly * 24 * 30:,}', GOLD)} {c('MYCEL/month', DIM)}")
    print(c("  Real earnings track utilisation, which is never 100%. See the", DIM))
    print(c("  live load per card at mycel.rent/#/rent.", DIM))
    print()
    print(c("  SUBSTRATE  ", DIM) + c("not accepting nodes yet", RED))
    print(c("  Settlement is not live, so this client cannot join anything and", DIM))
    print(c("  does not send your data anywhere. It reads your box and stops.", DIM))
    print(c("  Put your place in the queue at mycel.rent/#/operate", DIM))
    print()
    return 0


def main():
    ap = argparse.ArgumentParser(description="MYCEL node readiness check")
    ap.add_argument("--json", action="store_true", help="machine-readable output")
    ap.add_argument("--bench", action="store_true", help="measure fp16 throughput (needs torch)")
    ap.add_argument("--version", action="version", version=f"mycel-node {VERSION}")
    args = ap.parse_args()

    data = collect(bench=args.bench)
    if args.json:
        print(json.dumps(data, indent=2))
        return 0 if data["gpus"] else 1
    return report(data)


if __name__ == "__main__":
    sys.exit(main())