#!/usr/bin/env python3
"""Read the night back off the dashboards and say which memory fault it looks like.

Purpose: pull the hours around a failure out of Prometheus, pull the engine's own log out
    of the container runtime, and turn both into a written diagnosis: what the evidence
    shows, which of the known faults fits it, and which do not. It changes nothing and
    starts nothing; it only reads.
Platform: all (spark, strix, mac, nvidia). Pure Python, no dependencies.
Minimum memory: 8 GB, which is the service being diagnosed; this script needs almost none.
Assumes: the Prometheus from this part's lab, reachable and holding the window you care
    about, and optionally a container name whose log can be read with `docker logs`.
    A query whose metric your machine does not publish returns nothing, and the report
    says so rather than guessing. The verdict is a ranked shortlist, not a conclusion:
    it is where to look first, and the last section lists what would confirm each one.

Usage: python3 diagnose-oom.py --since 24h
       python3 diagnose-oom.py --since 12h --container local-gateway-swap-1
       python3 diagnose-oom.py --since 24h --output oom-report.md --labbook labbook.md
"""
import argparse
import json
import re
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

DURATION = re.compile(r"^(\d+)([smhd])$")
SECONDS = {"s": 1, "m": 60, "h": 3600, "d": 86400}

# Patterns that mean "the machine ran out of memory", in the words each stack uses. None of
# these is a guess: they are the strings the engines and the kernel print.
LOG_PATTERNS = [
    ("cuda-oom", re.compile(r"CUDA (?:error: )?out of memory|cudaErrorMemoryAllocation", re.I)),
    ("hip-oom", re.compile(r"hipErrorOutOfMemory|HIP out of memory", re.I)),
    ("metal-oom", re.compile(r"Insufficient Memory|failed to create (?:MTL|Metal) buffer", re.I)),
    ("host-oom", re.compile(r"std::bad_alloc|Cannot allocate memory|failed to allocate", re.I)),
    ("killed", re.compile(r"\bKilled\b|Out of memory: Kill(?:ed)? process|exit(?:ed)? .*137", re.I)),
    ("kv-cache", re.compile(r"KV cache|kv_cache|context size|n_ctx", re.I)),
    ("slot", re.compile(r"slot (?:is )?unavailable|no slot available|all slots are busy", re.I)),
]

# Every query is asked of every machine. A metric nothing publishes returns an empty result
# and is reported as absent, which is itself evidence about what you can and cannot see.
QUERIES = {
    "gateway_up": 'up{job="gateway"}',
    "engine_up": 'up{job=~"llama-server|vllm|sglang"}',
    "accelerator_memory_used": (
        'DCGM_FI_DEV_FB_USED * 1024 * 1024 or local_llm_accelerator_memory_used_bytes'
    ),
    "accelerator_memory_total": (
        '(DCGM_FI_DEV_FB_USED + DCGM_FI_DEV_FB_FREE) * 1024 * 1024 '
        'or local_llm_accelerator_memory_total_bytes'
    ),
    "kv_cache_usage": 'vllm:kv_cache_usage_perc or sglang:token_usage',
    "context_high_water": 'llamacpp:n_tokens_max',
    "requests_running": (
        'llamacpp:requests_processing or vllm:num_requests_running or sglang:num_running_reqs'
    ),
    "requests_waiting": (
        'llamacpp:requests_deferred or vllm:num_requests_waiting or sglang:num_queue_reqs'
    ),
    "host_memory_available": 'node_memory_MemAvailable_bytes',
}


def parse_duration(text):
    match = DURATION.match(text)
    if not match:
        sys.exit(f"diagnose-oom: --since takes a number and one of s, m, h, d, for example 24h; "
                 f"got {text!r}")
    return int(match.group(1)) * SECONDS[match.group(2)]


def query_range(base, expression, start, end, step):
    """One Prometheus range query. Returns a list of (labels, samples) or None on failure."""
    params = urllib.parse.urlencode({
        "query": expression, "start": f"{start:.3f}", "end": f"{end:.3f}", "step": step,
    })
    url = f"{base.rstrip('/')}/api/v1/query_range?{params}"
    try:
        with urllib.request.urlopen(url, timeout=30) as response:  # noqa: S310 - a URL you gave
            payload = json.loads(response.read().decode("utf-8"))
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, OSError) as exc:
        print(f"diagnose-oom: querying Prometheus failed: {exc}", file=sys.stderr)
        return None
    if payload.get("status") != "success":
        print(f"diagnose-oom: Prometheus refused the query: {payload.get('error')}", file=sys.stderr)
        return None
    return [(series.get("metric", {}),
             [(float(t), float(v)) for t, v in series.get("values", []) if v not in ("NaN",)])
            for series in payload["data"].get("result", [])]


def series_shape(samples):
    """Describe one series in the terms the diagnosis needs, without inventing precision."""
    values = [v for _, v in samples]
    if not values:
        return None
    first_third = values[: max(1, len(values) // 3)]
    last_third = values[-max(1, len(values) // 3):]
    start_level = sum(first_third) / len(first_third)
    end_level = sum(last_third) / len(last_third)
    peak = max(values)
    trough = min(values)
    rose = end_level > start_level * 1.15 and peak > trough * 1.15
    fell_at_end = values[-1] < peak * 0.5 < peak
    return {
        "samples": len(values),
        "start_level": start_level,
        "end_level": end_level,
        "peak": peak,
        "trough": trough,
        "rising": rose,
        "collapsed_at_end": fell_at_end,
        "ever_above_zero": peak > 0,
    }


def collect_metrics(base, start, end, step):
    out = {}
    for name, expression in QUERIES.items():
        result = query_range(base, expression, start, end, step)
        if result is None:
            out[name] = {"available": False, "reason": "Prometheus could not be queried"}
            continue
        if not result:
            out[name] = {"available": False, "reason": "no series; nothing publishes this here"}
            continue
        shapes = []
        for labels, samples in result:
            shape = series_shape(samples)
            if shape:
                shape["labels"] = labels
                shapes.append(shape)
        out[name] = {"available": bool(shapes), "series": shapes,
                     "reason": "" if shapes else "series present but empty over this window"}
    return out


def read_logs(container, since, runtime="docker"):
    if shutil.which(runtime) is None:
        return None, f"{runtime} is not on PATH"
    try:
        out = subprocess.run([runtime, "logs", "--since", since, container],
                             capture_output=True, text=True, timeout=60, check=False)
    except (OSError, subprocess.TimeoutExpired) as exc:
        return None, f"reading the log failed: {exc}"
    if out.returncode != 0:
        return None, f"{runtime} logs exited {out.returncode}: {out.stderr.strip()[:200]}"
    return (out.stdout or "") + (out.stderr or ""), ""


def scan_logs(text):
    hits = {}
    for name, pattern in LOG_PATTERNS:
        found = [line.strip() for line in text.splitlines() if pattern.search(line)]
        if found:
            hits[name] = found[-5:]
    return hits


def judge(metrics, log_hits):
    """Rank the four faults this part teaches against the evidence actually collected."""
    verdicts = []

    def shapes(name):
        block = metrics.get(name, {})
        return block.get("series", []) if block.get("available") else []

    def measured(name):
        """Absence of data is not evidence. Only reason from series that exist."""
        return bool(metrics.get(name, {}).get("available"))

    memory = shapes("accelerator_memory_used")
    kv = shapes("kv_cache_usage")
    context = shapes("context_high_water")
    waiting = shapes("requests_waiting")
    running = shapes("requests_running")

    memory_rising = any(s["rising"] for s in memory)
    kv_rising = any(s["rising"] for s in kv)
    context_rising = any(s["rising"] for s in context)
    was_queueing = any(s["ever_above_zero"] for s in waiting)
    running_peak = max((s["peak"] for s in running), default=0.0)
    memory_never_fell = memory and not any(s["collapsed_at_end"] for s in memory)

    # 1. Context creep: the cache fills because prompts got longer, not because there are
    #    more of them.
    score = 0
    why = []
    if kv_rising:
        score += 2
        why.append("key-value cache occupancy rose across the window")
    if context_rising:
        score += 2
        why.append("the high-water mark of context observed rose across the window")
    if measured("requests_waiting") and not was_queueing:
        score += 1
        why.append("nothing was queueing, so the load was not more requests")
    if "kv-cache" in log_hits:
        score += 1
        why.append("the log mentions the key-value cache or the context size")
    verdicts.append(("Context creep: the same traffic, with longer prompts", score, why))

    # 2. Concurrency beyond plan.
    score = 0
    why = []
    if was_queueing:
        score += 3
        why.append("requests were waiting for a slot")
    if running_peak > 1:
        score += 1
        why.append("more than one request was in flight at once")
    if "slot" in log_hits:
        score += 2
        why.append("the log says a slot was unavailable")
    verdicts.append(("Concurrency beyond what the engine was started for", score, why))

    # 3. Fragmentation or a leak: memory rises and never comes back down.
    score = 0
    why = []
    if memory_rising:
        score += 2
        why.append("accelerator memory in use rose across the window")
    if measured("accelerator_memory_used") and memory_never_fell:
        score += 2
        why.append("memory never returned to its starting level between bursts")
    if not kv_rising and not context_rising and memory_rising:
        score += 1
        why.append("memory rose while the cache and the context did not, so something else grew")
    verdicts.append(("Fragmentation or a leak: memory that is never given back", score, why))

    # 4. A model swap that never unloaded.
    score = 0
    why = []
    if memory_rising and memory_never_fell and measured("accelerator_memory_used"):
        score += 1
        why.append("memory stepped up and stayed up, which is what a load without an unload "
                   "looks like")
    engine = shapes("engine_up")
    if len(engine) > 1:
        score += 2
        why.append(f"{len(engine)} engine targets were up in this window, so more than one was running")
    verdicts.append(("A model swap that loaded without unloading", score, why))

    # 5. The plain one: something outside the engine took the memory.
    score = 0
    why = []
    if "host-oom" in log_hits or "killed" in log_hits:
        score += 3
        why.append("the log shows an allocation failure or a killed process")
    host = shapes("host_memory_available")
    if any(s["rising"] is False and s["end_level"] < s["start_level"] * 0.5 for s in host):
        score += 2
        why.append("host memory available fell by more than half across the window")
    verdicts.append(("The machine ran out, not the engine: something else took the memory", score, why))

    for name in ("cuda-oom", "hip-oom", "metal-oom"):
        if name in log_hits:
            verdicts.append((f"Confirmed by the log: an allocation on the accelerator failed ({name})",
                             5, [f"the log contains: {log_hits[name][-1][:160]}"]))

    verdicts.sort(key=lambda row: row[1], reverse=True)
    return verdicts


def human_bytes(value):
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if abs(value) < 1024 or unit == "TiB":
            return f"{value:,.1f} {unit}"
        value /= 1024
    return f"{value:,.1f} TiB"


def write_report(metrics, log_hits, log_note, verdicts, window, out):
    lines = ["# Memory failure: what the record shows", ""]
    lines.append(f"Window examined: {window}")
    lines.append(f"Written: {datetime.now(timezone.utc).isoformat(timespec='seconds')}")
    lines.append("")
    lines.append("## What was measurable")
    lines.append("")
    for name, block in metrics.items():
        if not block.get("available"):
            lines.append(f"- **{name}**: nothing to read ({block.get('reason')})")
            continue
        for shape in block["series"]:
            label = ", ".join(f"{k}={v}" for k, v in shape["labels"].items()
                              if k not in ("__name__",)) or "no labels"
            fmt = human_bytes if "memory" in name else (lambda v: f"{v:,.3f}")
            lines.append(
                f"- **{name}** [{label}]: started near {fmt(shape['start_level'])}, "
                f"ended near {fmt(shape['end_level'])}, peaked at {fmt(shape['peak'])}"
                f"{'; rising across the window' if shape['rising'] else ''}"
                f"{'; collapsed at the end' if shape['collapsed_at_end'] else ''}")
    lines.append("")
    lines.append("## What the log said")
    lines.append("")
    if log_note:
        lines.append(f"- No log was read: {log_note}")
    elif not log_hits:
        lines.append("- The log was read and none of the known memory-failure patterns appeared. "
                     "That is worth knowing: it moves the question towards the host and away "
                     "from the engine.")
    else:
        for name, found in log_hits.items():
            lines.append(f"- **{name}**, most recent occurrence:")
            for line in found[-2:]:
                lines.append(f"  - `{line[:200]}`")
    lines.append("")
    lines.append("## Where to look first")
    lines.append("")
    lines.append("Ranked by how much of the evidence fits, not by how likely each is in general. "
                 "A candidate with no supporting evidence is listed so that you can see it was "
                 "considered and found wanting.")
    lines.append("")
    for name, score, why in verdicts:
        lines.append(f"### {name}")
        lines.append("")
        lines.append(f"Evidence for it: {score if score else 'none'}")
        lines.append("")
        if why:
            for reason in why:
                lines.append(f"- {reason}")
        else:
            lines.append("- Nothing in the collected evidence points here.")
        lines.append("")
    lines.append("## What would settle it")
    lines.append("")
    lines.extend([
        "- **Context creep**: compare the average prompt length this week with last week. "
        "If it doubled, the cache filling is arithmetic rather than a fault.",
        "- **Concurrency**: count the distinct callers in the gateway's usage records over the "
        "hour before the failure, and compare with the slot count the engine was started with.",
        "- **Fragmentation or a leak**: restart the engine with the same load and watch whether "
        "memory returns to the same starting level. If it does not, it is the process.",
        "- **A swap that never unloaded**: ask the model manager what it thinks is loaded, and "
        "compare with what the accelerator says is resident. Those two disagreeing is the fault.",
        "- **The machine, not the engine**: look at host memory and at what else was running. "
        "The engine is often the largest process rather than the guilty one.",
        "",
        "Then change one thing, put the same load back, and watch the same graph. A fix that "
        "cannot be seen on the graph that showed the fault has not been proved.",
    ])
    text = "\n".join(lines) + "\n"
    if out:
        Path(out).expanduser().write_text(text, encoding="utf-8")
        print(f"wrote {out}")
    else:
        print(text)
    return text


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--prometheus", default="http://127.0.0.1:9090",
                        help="base URL of your Prometheus")
    parser.add_argument("--since", default="24h", help="how far back to look, e.g. 24h or 3d")
    parser.add_argument("--step", default="5m", help="resolution of the range queries")
    parser.add_argument("--container", default=None, help="container whose log to read")
    parser.add_argument("--runtime", default="docker", help="docker or podman")
    parser.add_argument("--output", default=None, help="write the report here")
    parser.add_argument("--labbook", default=None, help="append one JSON line here")
    args = parser.parse_args()

    seconds = parse_duration(args.since)
    end = time.time()
    start = end - seconds
    window = (f"the last {args.since}, from "
              f"{datetime.fromtimestamp(start, timezone.utc).isoformat(timespec='seconds')} "
              f"to {datetime.fromtimestamp(end, timezone.utc).isoformat(timespec='seconds')}")

    metrics = collect_metrics(args.prometheus, start, end, args.step)
    log_hits, log_note = {}, "no --container was given"
    if args.container:
        text, note = read_logs(args.container, args.since, args.runtime)
        if text is None:
            log_note = note
        else:
            log_hits, log_note = scan_logs(text), ""

    verdicts = judge(metrics, log_hits)
    write_report(metrics, log_hits, log_note, verdicts, window, args.output)

    if args.labbook:
        top = verdicts[0] if verdicts else ("no candidate", 0, [])
        record = {
            "lab": "part-23/challenge-oom",
            "window": args.since,
            "top_candidate": top[0],
            "evidence_score": top[1],
            "log_patterns_seen": sorted(log_hits.keys()),
            "metrics_unavailable": sorted(k for k, v in metrics.items() if not v.get("available")),
            "recorded": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        }
        with Path(args.labbook).expanduser().open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record, sort_keys=True) + "\n")
        print(f"recorded in {args.labbook}")


if __name__ == "__main__":
    main()