#!/usr/bin/env python3
"""parse_llama_log.py — normalize a -lv 4 llama-server startup log into a
compact placement record (JSON).

The local model must never read raw llama.cpp startup logs. This script
extracts the hard numbers (per-device model/KV/RS/compute/finish buffers,
fitted ctx, graph splits, load status) so the supervisor model works on a
~20-line record instead of 500 lines of log.

Usage:
  parse_llama_log.py LOG [options]

Options:
  --rocm-smi SNAPSHOT   file from `rocm-smi --showmeminfo vram` (or '-' = run live)
  --device-order "rocm0,rocm1"   map GPU[i] -> device name (default rocm0,rocm1)
  --fit-target N        record fit target (not always in the log)
  --tensor-split S      record tensor split (from launcher)
  --ot REGEX            record the -ot override used (from launcher)
  --key-json '{...}'    placement-key object to attach (model, caches, build, ...)
  --append HISTORY.jsonl   append {key, run} to the history file
  --quiet               only print the final JSON

Exit: 0 = parsed (check load_ok in output), 2 = input error.
"""
import argparse
import json
import re
import subprocess
import sys
import time

FAIL_PATTERNS = [
    (r"allocation failed", "allocation_failed"),
    (r"out of memory", "oom"),
    (r"cannot allocate", "allocation_failed"),
    (r"GGML_ASSERT", "assert"),
    # "failed to fit ... n_gpu_layers already set by user" is a benign fork
    # warning: the ctx reduction is still applied, only n_gpu_layers is kept.
    (r"failed to fit(?! params to free device memory: n_gpu_layers)|fit failed", "fit_failed"),
    (r"error:.*device", "device_error"),
    (r"segfault|segmentation fault", "crash"),
]


def f2(x):
    return round(float(x), 2)


def add(d, k, v):
    d[k] = f2(d.get(k, 0.0) + float(v))


def parse_rocm_smi(path):
    """Return list of {index, total_mib, used_mib, free_mib}."""
    if path == "-":
        out = subprocess.run(
            ["rocm-smi", "--showmeminfo", "vram"],
            capture_output=True, text=True, timeout=30,
        ).stdout
    else:
        out = open(path).read()
    per = {}
    for m in re.finditer(r"GPU\[(\d+)\]\s*:\s*VRAM (Total Used|Total) Memory \(B\):\s*(\d+)", out):
        idx, kind, val = int(m.group(1)), m.group(2), int(m.group(3)) / (1024 * 1024)
        per.setdefault(idx, {})["used" if "Used" in kind else "total"] = f2(val)
    res = []
    for idx in sorted(per):
        t, u = per[idx].get("total"), per[idx].get("used")
        res.append({
            "index": idx,
            "total_mib": t,
            "used_mib": u,
            "free_mib": f2(t - u) if (t is not None and u is not None) else None,
        })
    return res


def parse_log(path):
    lines = open(path, errors="replace").read().splitlines()
    text = "\n".join(lines)
    run = {
        "schema": 1,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "build": None,
        "model": None,
        "n_layer": None,
        "devices": {},      # name -> {total_mib, free_at_load_mib}
        "model_mib": {},    # dev -> target model buffer
        "kv_mib": {},       # dev -> all KV buffers (target + draft)
        "kv_draft_mib": 0.0,  # total draft (MTP) KV across devices
        "rs_mib": {},       # dev -> recurrent state
        "compute_mib": {},  # dev -> sum of all sched compute buffers (may share)
        "finish_mib": {},   # dev -> recurrent checkpoint (finish) buffer
        "used_log_mib": {}, # dev -> model+kv+rs+compute+finish (log-derived)
        "free_mib": {},     # dev -> rocm-smi free (authoritative when present)
        "graph_splits": [],
        "sched_copies": [],
        "ctx_fitted": None,
        "ctx_reduced_from": None,
        "ctx_reduced_to": None,
        "cache_target": None,   # "q8_0/q5_1"
        "cache_draft": None,    # "q4_0/q4_0"
        "listening": False,
        "listen_url": None,
        "failures": [],
        "load_ok": False,
    }

    m = re.search(r"build (\d+) \(([\w.-]+)\)", text)
    if m:
        run["build"] = f"{m.group(1)} ({m.group(2)})"
    m = re.search(r"load_model: loading model '([^']+)'", text)
    if m:
        run["model"] = m.group(1)
    m = re.search(r"n_layer\s+=\s+(\d+)", text)
    if m:
        run["n_layer"] = int(m.group(1))

    for m in re.finditer(r"-\s+(ROCm\d+|Vulkan\d+|CUDA\d+|CPU)\s*:.*?\((\d+) MiB, (\d+) MiB free\)", text):
        run["devices"][m.group(1)] = {
            "total_mib": int(m.group(2)),
            "free_at_load_mib": int(m.group(3)),
        }

    for m in re.finditer(r"load_tensors:\s+(\S+)\s+model buffer size =\s+([\d.]+) MiB", text):
        add(run["model_mib"], m.group(1), m.group(2))
    # context order: 1st constructing = target, 2nd = MTP draft
    nctx = 0
    for line in lines:
        if "constructing llama_context" in line:
            nctx += 1
        m = re.search(r"llama_kv_cache:\s+(\S+)\s+KV buffer size =\s+([\d.]+) MiB", line)
        if m:
            add(run["kv_mib"], m.group(1), m.group(2))
            if nctx >= 2:
                run["kv_draft_mib"] = f2(run["kv_draft_mib"] + float(m.group(2)))
    for m in re.finditer(r"llama_memory_recurrent:\s+(\S+)\s+RS buffer size =\s+([\d.]+) MiB", text):
        add(run["rs_mib"], m.group(1), m.group(2))
    for m in re.finditer(r"sched_reserve:\s+(\S+)\s+compute buffer size =\s+([\d.]+) MiB", text):
        add(run["compute_mib"], m.group(1), m.group(2))
    for m in re.finditer(r"finish: allocated '([^']+)' buffer ([\d.]+) MiB", text):
        add(run["finish_mib"], m.group(1), m.group(2))

    run["graph_splits"] = [int(x) for x in re.findall(r"sched_reserve: graph splits = (\d+)", text)]
    run["sched_copies"] = [int(x) for x in re.findall(r"sched copies = (\d+)", text)]

    ctxs = [int(x) for x in re.findall(r"n_ctx_slot = (\d+)", text)]
    if ctxs:
        run["ctx_fitted"] = max(ctxs)
    # fitter's internal projection at initial params + per-device free targets
    run["fit_proj_mib"] = {}
    run["fit_target_free_mib"] = {}
    for m in re.finditer(
            r"common_params_fit_impl:\s+- (\S+) \(.*?\)\s*:\s+(\d+) total,\s+(\d+) used,\s+-?\d+ free vs\. target of\s+(\d+)",
            text):
        dev = m.group(1)
        if dev not in run["devices"] and dev != "CPU":
            continue
        run["fit_proj_mib"][dev] = int(m.group(3))
        run["fit_target_free_mib"][dev] = int(m.group(4))

    m = re.search(r"context size reduced from (\d+) to (\d+)", text)
    if m:
        run["ctx_reduced_from"] = int(m.group(1))
        run["ctx_reduced_to"] = int(m.group(2))

    # the device scan lists every compiled backend (e.g. ROCm0/1 lines even
    # on a Vulkan run); keep only devices the fitter actually projected
    if run["fit_proj_mib"]:
        keep = set(run["fit_proj_mib"]) | {"CPU"}
        run["devices"] = {d: v for d, v in run["devices"].items() if d in keep}

    m = re.search(r"K \((\w+)\):.*?V \((\w+)\):", text)
    if m:
        run["cache_target"] = f"{m.group(1)}/{m.group(2)}"
    m = re.search(r"cache_k=(\w+), cache_v=(\w+)", text)
    if m:
        run["cache_draft"] = f"{m.group(1)}/{m.group(2)}"

    m = re.search(r"listening on (http://\S+)", text)
    run["listening"] = bool(m)
    if m:
        run["listen_url"] = m.group(1)

    for pat, label in FAIL_PATTERNS:
        if re.search(pat, text, re.IGNORECASE):
            run["failures"].append(label)

    # used derived from log (approx: compute buffers may be shared across
    # contexts; rocm-smi free_mib is authoritative when available)
    devs = set(run["model_mib"]) | set(run["kv_mib"]) | set(run["rs_mib"]) | set(run["compute_mib"]) | set(run["finish_mib"])
    for d in sorted(devs):
        run["used_log_mib"][d] = f2(
            run["model_mib"].get(d, 0) + run["kv_mib"].get(d, 0) +
            run["rs_mib"].get(d, 0) + run["compute_mib"].get(d, 0) +
            run["finish_mib"].get(d, 0),
        )

    run["load_ok"] = run["listening"] and not run["failures"] and run["ctx_fitted"] is not None
    return run


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("log")
    ap.add_argument("--rocm-smi", default=None, help="snapshot file, or '-' for live")
    ap.add_argument("--device-order", default="rocm0,rocm1")
    ap.add_argument("--fit-target", type=int, default=None)
    ap.add_argument("--tensor-split", default=None)
    ap.add_argument("--ot", default=None)
    ap.add_argument("--key-json", default=None)
    ap.add_argument("--append", default=None, metavar="HISTORY.jsonl")
    ap.add_argument("--quiet", action="store_true")
    a = ap.parse_args()

    try:
        run = parse_log(a.log)
    except OSError as e:
        print(f"error: cannot read log: {e}", file=sys.stderr)
        sys.exit(2)

    if a.rocm_smi:
        try:
            smi = parse_rocm_smi(a.rocm_smi)
        except Exception as e:
            print(f"warning: rocm-smi parse failed: {e}", file=sys.stderr)
            smi = []
        order = [x.strip() for x in a.device_order.split(",") if x.strip()]
        # match against the log's device naming (e.g. rocm0 -> ROCm0)
        known = list(run["model_mib"]) + list(run["devices"])
        def canon(name):
            for k in known:
                if k.lower() == name.lower():
                    return k
            return name
        for g in smi:
            name = canon(order[g["index"]]) if g["index"] < len(order) else f"GPU{g['index']}"
            run["free_mib"][name] = g["free_mib"]
            run.setdefault("rocm_smi", {})[name] = g

    if a.fit_target is not None:
        run["fit_target"] = a.fit_target
    if a.tensor_split:
        run["tensor_split"] = a.tensor_split
    if a.ot:
        run["ot"] = a.ot

    if a.append:
        key = json.loads(a.key_json) if a.key_json else {}
        if not key and run.get("model"):
            key = {"model_file": run["model"].rsplit("/", 1)[-1]}
        rec = {"schema": 1, "key": key, "run": run}
        with open(a.append, "a") as f:
            f.write(json.dumps(rec) + "\n")
        if not a.quiet:
            print(f"appended to {a.append}")

    print(json.dumps({"key": (json.loads(a.key_json) if a.key_json else None), "run": run}, indent=2))
    return 0


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