#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
"""collect.py - aggregate perf_workload PERF lines and host vgsp evidence.

Three input classes, all optional in one invocation:

  --guest RUN.txt [RUN.txt ...]   perf_workload stdout captures (one per run)
  --host-gr DMESG.txt             dmesg capture (raw monotonic timestamps
                                  preferred; dmesg -T is accepted but only the
                                  per-launch elapsed_ms is then usable)
  --state-before / --state-after  /sys/kernel/debug/vgsp/gsp-phys-state snaps

Output is an `open-iov-perf/1` JSON document (stdout or --out FILE).

Compare mode:

  collect.py --compare HOST.json GUEST.json --out compare.json \
             [--thresholds thresholds.json]

The host artifact is the bare-metal oracle; every ratio is guest/host, so
"lower is better" metrics are good below 1.0 and "higher is better" metrics
are good above 1.0. Exit status is 0 only when both sides verified
(bad=0, STATUS OK) and all supplied thresholds pass.
"""

import argparse
import json
import re
import statistics
import sys
from datetime import datetime, timezone

SCHEMA = "open-iov-perf/1"

RAW_TS = re.compile(r"^\[\s*(\d+\.\d+)\]\s*(.*)$")
T_TS = re.compile(r"^\[(\w{3} \w{3}\s+\d+ \d{2}:\d{2}:\d{2} \d{4})\]\s*(.*)$")
GR_OK = re.compile(r"GUEST GR LAUNCH OK .*?elapsed_ms=(\d+)\s+ok=(\d+)")
GR_SUBMIT = re.compile(r"GR LAUNCH submitted")


def parse_perf_file(path):
    records = []
    with open(path, "r", errors="replace") as fh:
        for line in fh:
            line = line.rstrip("\n").rstrip("\r")
            if not line.startswith("PERF "):
                continue
            rest = line[5:].strip()
            if not rest:
                continue
            parts = rest.split(None, 1)
            fields = {}
            kind = parts[0]
            if "=" in parts[0]:
                k, v = parts[0].split("=", 1)
                kind = k
                fields[k] = v
            if len(parts) > 1:
                for tok in parts[1].split():
                    if "=" in tok:
                        k, v = tok.split("=", 1)
                        fields[k] = v
                if len(parts[1].split()) == 1 and "=" not in parts[1]:
                    fields["value"] = parts[1]
            records.append({"kind": kind, "fields": fields})
    return records


def fnum(v, default=None):
    try:
        return float(v)
    except (TypeError, ValueError):
        return default


def inum(v, default=None):
    try:
        return int(v, 0)
    except (TypeError, ValueError):
        return default


def stats_of(values):
    if not values:
        return None
    vals = sorted(values)
    n = len(vals)

    def pct(p):
        return vals[min(n - 1, int(p / 100.0 * (n - 1) + 0.5))]

    return {
        "n": n,
        "values": values,
        "min": vals[0],
        "max": vals[-1],
        "mean": statistics.fmean(values),
        "median": statistics.median(values),
        "p50": pct(50),
        "p95": pct(95),
    }


def extract_run(records):
    """Single perf_workload run -> flat metric dict + verification."""
    m = {}
    ver = {"status": None, "bad": None}
    for rec in records:
        k = rec["kind"]
        f = rec["fields"]
        if k == "device":
            m["device"] = f.get("device", "")
        elif k == "run_cfg":
            m["run_cfg"] = f
        elif k in ("init_ms", "malloc_first_ms", "launch_first_ms"):
            v = fnum(f.get(k))
            if v is not None:
                m[k] = v
        elif k == "malloc_ms":
            m["malloc_ms_p50"] = fnum(f.get("p50"))
            m["malloc_ms_p95"] = fnum(f.get("p95"))
            m["malloc_ms_mean"] = fnum(f.get("mean"))
        elif k == "launch_warm":
            m["launch_warm_p50"] = fnum(f.get("p50"))
            m["launch_warm_p95"] = fnum(f.get("p95"))
            m["launch_warm_mean"] = fnum(f.get("mean"))
        elif k == "kernel_fma":
            m["kernel_fma_ms"] = fnum(f.get("ms"))
            m["kernel_fma_gflops"] = fnum(f.get("gflops"))
        elif k == "throughput":
            m["throughput_kps"] = fnum(f.get("kps"))
            m["throughput_wall_ms"] = fnum(f.get("wall_ms"))
        elif k == "copy":
            d = f.get("dir", "?")
            b = f.get("bytes", "?")
            m["copy_%s_%s_gbps" % (d, b)] = fnum(f.get("gbps"))
            m["copy_%s_%s_ms" % (d, b)] = fnum(f.get("ms"))
        elif k == "clock64":
            m["clock64_available"] = inum(f.get("available"))
            m["clock64_min"] = fnum(f.get("min"))
            m["clock64_median"] = fnum(f.get("median"))
        elif k == "checksum":
            ver["bad"] = inum(f.get("bad"))
        elif k == "STATUS":
            ver["status"] = f.get("value") or f.get("STATUS") or f.get("status")
    return m, ver


def parse_dmesg(path, since=None):
    launches = []
    pending_submit = None
    with open(path, "r", errors="replace") as fh:
        for line in fh:
            ts = None
            body = line
            mo = RAW_TS.match(line)
            if mo:
                ts = float(mo.group(1))
                body = mo.group(2)
            else:
                mo = T_TS.match(line)
                if mo:
                    ts = mo.group(1)
                    body = mo.group(2)
            if isinstance(ts, float) and since is not None and ts <= since:
                continue
            if GR_SUBMIT.search(body):
                pending_submit = ts
            mo = GR_OK.search(body)
            if mo:
                entry = {
                    "elapsed_ms": int(mo.group(1)),
                    "ok": int(mo.group(2)),
                }
                if isinstance(ts, float) and isinstance(pending_submit,
                                                         float):
                    entry["submit_to_ok_ms"] = round(
                        (ts - pending_submit) * 1000.0, 3)
                launches.append(entry)
                pending_submit = None
    out = {"launches": launches}
    for key in ("elapsed_ms", "submit_to_ok_ms"):
        vals = [e[key] for e in launches if key in e]
        if vals:
            out[key] = stats_of(vals)
    return out


def parse_state(path):
    sections = {}
    if not path:
        return sections
    with open(path, "r", errors="replace") as fh:
        for line in fh:
            line = line.strip()
            mo = re.match(r"^([A-Za-z0-9_\-]+):\s*(.*)$", line)
            if not mo:
                continue
            fields = {}
            for tok in mo.group(2).split():
                if "=" in tok:
                    k, v = tok.split("=", 1)
                    iv = inum(v)
                    fields[k] = iv if iv is not None else v
            if fields:
                sections[mo.group(1)] = fields
    return sections


def state_delta(before, after):
    delta = {}
    for sec, fields in after.items():
        if sec not in before:
            continue
        d = {}
        for k, v in fields.items():
            bv = before[sec].get(k)
            if isinstance(v, int) and isinstance(bv, int):
                d[k] = v - bv
        if d:
            delta[sec] = d
    return delta


def parse_meta(path):
    env = {}
    if not path:
        return env
    with open(path, "r", errors="replace") as fh:
        for line in fh:
            line = line.strip()
            mo = re.match(r"^([a-z0-9_]+):\s*(.*)$", line)
            if mo:
                env[mo.group(1)] = mo.group(2)
    return env


def parse_marker(path):
    if not path:
        return None
    with open(path, "r", errors="replace") as fh:
        for line in fh:
            mo = RAW_TS.match(line)
            if mo:
                return float(mo.group(1))
    return None


def build_document(args):
    metric_runs = []
    verification = {"status": None, "bad": None}
    device = None
    run_cfg = None
    for path in args.guest:
        records = parse_perf_file(path)
        m, ver = extract_run(records)
        metric_runs.append(m)
        device = device or m.get("device")
        run_cfg = run_cfg or m.get("run_cfg")
        if ver["status"] is not None:
            verification["status"] = ver["status"]
        if ver["bad"] is not None:
            verification["bad"] = ver["bad"]

    keys = sorted({k for m in metric_runs for k in m})
    aggregate = {}
    for k in keys:
        if k in ("device", "run_cfg"):
            continue
        vals = [m[k] for m in metric_runs if isinstance(m.get(k), (int, float))]
        s = stats_of(vals)
        if s:
            aggregate[k] = s

    doc = {
        "schema": SCHEMA,
        "created": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "role": args.role,
        "runs": len(args.guest),
        "device": device,
        "run_cfg": run_cfg,
        "verification": verification,
        "per_run": metric_runs,
        "aggregate": aggregate,
    }
    if args.host_gr:
        doc["host_gr"] = parse_dmesg(args.host_gr,
                                     since=parse_marker(args.host_gr_marker))
        if args.host_gr_marker:
            doc["host_gr"]["since"] = parse_marker(args.host_gr_marker)
    if args.state_before or args.state_after:
        before = parse_state(args.state_before)
        after = parse_state(args.state_after)
        doc["counters"] = {
            "before": before,
            "after": after,
            "delta": state_delta(before, after),
        }
    if args.meta:
        doc["environment"] = parse_meta(args.meta)
    return doc


LOWER_BETTER = {
    "init_ms", "malloc_first_ms", "malloc_ms_p50", "malloc_ms_p95",
    "malloc_ms_mean", "launch_first_ms", "launch_warm_p50",
    "launch_warm_p95", "launch_warm_mean", "kernel_fma_ms", "clock64_min",
    "clock64_median",
}
NO_COMPARE = {"clock64_available"}


def lower_is_better(key):
    return key in LOWER_BETTER or key.endswith("_ms")


def metric_view(doc):
    view = {}
    for k, s in doc.get("aggregate", {}).items():
        view[k] = s.get("p50", s.get("mean"))
    view["host_gr_elapsed_ms"] = doc.get("host_gr", {}).get(
        "elapsed_ms", {}).get("p50")
    view["host_gr_submit_to_ok_ms"] = doc.get("host_gr", {}).get(
        "submit_to_ok_ms", {}).get("p50")
    view["bad"] = doc.get("verification", {}).get("bad")
    return {k: v for k, v in view.items() if v is not None}


def run_compare(args):
    with open(args.compare[0]) as fh:
        host = json.load(fh)
    with open(args.compare[1]) as fh:
        guest = json.load(fh)
    thresholds = {}
    if args.thresholds:
        with open(args.thresholds) as fh:
            thresholds = json.load(fh).get("metrics", {})

    hm = metric_view(host)
    gm = metric_view(guest)
    rows = []
    failed = False
    for key in sorted(set(hm) | set(gm)):
        if key == "bad" or key in NO_COMPARE:
            continue
        h, g = hm.get(key), gm.get(key)
        row = {"metric": key, "host": h, "guest": g}
        if isinstance(h, (int, float)) and isinstance(g, (int, float)):
            row["guest_over_host"] = g / h if h else None
            row["better"] = "lower" if lower_is_better(key) else "higher"
            rule = thresholds.get(key)
            if rule:
                limit = rule.get("limit")
                mode = rule.get("mode", "max_ratio")
                ratio = row["guest_over_host"]
                if mode == "max_ratio":
                    ok = ratio is not None and ratio <= limit
                else:
                    ok = ratio is not None and ratio >= limit
                row["limit"] = limit
                row["verdict"] = "PASS" if ok else "FAIL"
                if not ok:
                    failed = True
        rows.append(row)

    host_ok = (host.get("verification", {}).get("bad") == 0 and
               host.get("verification", {}).get("status") == "OK")
    guest_ok = (guest.get("verification", {}).get("bad") == 0 and
                guest.get("verification", {}).get("status") == "OK")
    out = {
        "schema": SCHEMA + "/compare",
        "created": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "host_verified": host_ok,
        "guest_verified": guest_ok,
        "thresholds": thresholds,
        "rows": rows,
    }
    if args.out:
        with open(args.out, "w") as fh:
            json.dump(out, fh, indent=2, sort_keys=True)
    width = max([len(r["metric"]) for r in rows] + [24])
    print("%-*s %12s %12s %10s %8s %6s" %
          (width, "metric", "host", "guest", "g/h", "better", "verdict"))
    for r in rows:
        h = r.get("host")
        g = r.get("guest")
        ratio = r.get("guest_over_host")
        print("%-*s %12s %12s %10s %8s %6s" % (
            width, r["metric"],
            "%.4g" % h if isinstance(h, (int, float)) else "-",
            "%.4g" % g if isinstance(g, (int, float)) else "-",
            "%.4g" % ratio if isinstance(ratio, (int, float)) else "-",
            r.get("better", ""), r.get("verdict", "")))
    print("verification: host=%s guest=%s" %
          ("OK" if host_ok else "FAIL", "OK" if guest_ok else "FAIL"))
    return 0 if (host_ok and guest_ok and not failed) else 1


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--guest", nargs="*", default=[])
    ap.add_argument("--host-gr")
    ap.add_argument("--host-gr-marker")
    ap.add_argument("--state-before")
    ap.add_argument("--state-after")
    ap.add_argument("--meta")
    ap.add_argument("--role", choices=["host", "guest"], default="guest")
    ap.add_argument("--out")
    ap.add_argument("--compare", nargs=2, metavar=("HOST.json", "GUEST.json"))
    ap.add_argument("--thresholds")
    args = ap.parse_args()

    if args.compare:
        return run_compare(args)

    doc = build_document(args)
    text = json.dumps(doc, indent=2, sort_keys=True)
    if args.out:
        with open(args.out, "w") as fh:
            fh.write(text + "\n")
        print("wrote %s (%d run(s), %d metric(s))" %
              (args.out, doc["runs"], len(doc["aggregate"])))
    else:
        print(text)
    return 0


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