#!/usr/bin/env python3
"""Measure the server under a realistic load with spark (roadmap Phase 9).

Boots the test server on the island world template once per scenario, adds the fake players and bases of
`/cobblesurvival perf` (mod: selftest/PerfLoad.kt), lets things settle and records every tick, while spark profiles the
server thread. Writes build/perf/<scenario>.json and .sparkprofile (open that in spark's viewer if you want to look closer) and docs/performance.md.

What it can't show: fake clients send no packets and encode none, so network cost is left out, and no battles are fought.

Usage: tools/perf_test.py [--scenarios solo group stress] [--settle 120] [--measure 180] [--no-doc]
"""
import argparse
import gzip
import json
import re
import shutil
import sys
import time
from collections import defaultdict
from pathlib import Path

from mcserver import BUILD, ROOT, SELFTEST_JVM_ARGS, SERVER, Server, eula_accepted, prepare, scan_log

TEMPLATE = BUILD / "world-template"
OUT = BUILD / "perf"
# players, pastures each, Pokémon per pasture, machines each
SCENARIOS = {"solo": (1, 0, 0, 0), "group": (5, 4, 8, 5), "stress": (8, 6, 16, 10)}
HEAP = "6G"


def fields(buf):
    """Protobuf wire format: (field number, wire type, value)."""
    i, n = 0, len(buf)

    def varint():
        nonlocal i
        shift = result = 0
        while True:
            b = buf[i]
            i += 1
            result |= (b & 0x7f) << shift
            if not b & 0x80:
                return result
            shift += 7

    while i < n:
        key = varint()
        number, wire = key >> 3, key & 7
        if wire == 0:
            yield number, wire, varint()
        elif wire == 1:
            yield number, wire, buf[i:i + 8]
            i += 8
        elif wire == 2:
            length = varint()
            yield number, wire, buf[i:i + length]
            i += length
        elif wire == 5:
            yield number, wire, buf[i:i + 4]
            i += 4
        else:
            raise ValueError(f"wire type {wire}")


def doubles(value):
    import struct
    return struct.unpack(f"<{len(value) // 8}d", value)


def varints(value):
    return [v for v in iter_varints(value)]


def iter_varints(value):
    i, n = 0, len(value)
    while i < n:
        shift = result = 0
        while True:
            b = value[i]
            i += 1
            result |= (b & 0x7f) << shift
            if not b & 0x80:
                break
            shift += 7
        yield result


def read_profile(path):
    """spark's sampler data: {class.method: self ms}, inclusive per method, total ms of the server thread."""
    raw = path.read_bytes()
    if raw[:2] == b"\x1f\x8b":
        raw = gzip.decompress(raw)
    nodes, roots = [], []
    for number, _, value in fields(raw):
        if number != 2:
            continue
        flat, children, time_ms = [], [], 0.0
        for n, _, v in fields(value):
            if n == 3:  # children: every stack frame, flattened
                node = {"class": "", "method": "", "time": 0.0, "kids": []}
                for fn, wire, fv in fields(v):
                    if fn == 3:
                        node["class"] = fv.decode()
                    elif fn == 4:
                        node["method"] = fv.decode()
                    elif fn == 8:
                        node["time"] += sum(doubles(fv)) if wire == 2 else doubles(fv)[0]
                    elif fn == 9:
                        node["kids"] += varints(fv) if wire == 2 else [fv]
                flat.append(node)
            elif n == 4:
                time_ms += sum(doubles(v))
            elif n == 5:
                children += varints(v)
        nodes.append(flat)
        roots.append((time_ms, children))
    self_ms, inclusive = defaultdict(float), defaultdict(float)
    total = 0.0
    for flat, (time_ms, children) in zip(nodes, roots):
        total += time_ms
        def walk(index, path_seen):
            node = flat[index]
            name = f"{node['class']}.{node['method']}"
            self_ms[name] += node["time"] - sum(flat[k]["time"] for k in node["kids"])
            fresh = name not in path_seen
            if fresh:
                inclusive[name] += node["time"]
                path_seen.add(name)
            for k in node["kids"]:
                walk(k, path_seen)
            if fresh:
                path_seen.discard(name)
        sys.setrecursionlimit(20000)
        for c in children:
            walk(c, set())
    return self_ms, inclusive, total


def group_of(name):
    if name.startswith("libc.so") or ".park" in name or "waitUntilNextTick" in name or "managedBlock" in name:
        return "Waiting for the next tick"
    if name.startswith("dev.cobblesurvival"):
        return "Companion mod"
    if name.startswith(("com.cobblemon", "com.cobbledex")):
        return "Cobblemon"
    if name.startswith(("net.minecraft", "com.mojang", "it.unimi")):
        return "Minecraft"
    if name.startswith(("java.", "jdk.", "sun.", "kotlin.")):
        return "JVM and Kotlin"
    return "Other mods"


def summarise(profile):
    self_ms, inclusive, total = read_profile(profile)
    groups = defaultdict(float)
    for name, ms in self_ms.items():
        groups[group_of(name)] += ms
    pick = lambda prefix: sorted(((n, ms) for n, ms in inclusive.items() if n.startswith(prefix)), key=lambda kv: -kv[1])[:8]
    return {"total_ms": total, "groups": dict(sorted(groups.items(), key=lambda kv: -kv[1])),
            "companion_top": pick("dev.cobblesurvival"), "cobblemon_top": pick("com.cobblemon"),
            "self_top": sorted(self_ms.items(), key=lambda kv: -kv[1])[:10]}


def parse_report(line):
    return {k: float(v) for k, v in re.findall(r"(\w+)=([\d.]+)", line)}


def run(name, config, args_file, recipe, settle, measure):
    players, pastures, pokemon, machines = config
    world = SERVER / "world"
    shutil.rmtree(world, ignore_errors=True)
    shutil.copytree(TEMPLATE, world)
    for old in SERVER.glob("config/spark/**/*.sparkprofile"):
        old.unlink()
    with Server(args_file, {"level-type": recipe["level_type"], "level-seed": recipe["seed"]}, heap=HEAP,
                stdout_name=f"perf-{name}-stdout.log", jvm_args=SELFTEST_JVM_ARGS) as server:
        boot = server.wait_ready()
        print(f"[{name}] booted in {boot:.0f}s: {server.command(f'cobblesurvival perf start {players} {pastures} {pokemon} {machines}').strip()}", flush=True)
        time.sleep(settle)
        server.command("cobblesurvival perf measure")
        print(f"[{name}] {server.command(f'spark profiler start --timeout {measure} --save-to-file').strip()[:120]}", flush=True)
        time.sleep(measure + 15)
        report = server.command("cobblesurvival perf report").strip()
        tps = re.sub(r"§.", "", server.command("spark tps"))
        health = re.sub(r"§.", "", server.command("spark health"))
        print(f"[{name}] {report}", flush=True)
        for _ in range(30):
            profiles = sorted(SERVER.glob("config/spark/**/*.sparkprofile"))
            if profiles:
                break
            time.sleep(2)
        server.stop()
    errors = [line[:200] for line in scan_log(server.log)][:30]
    result = {"scenario": name, "players": players, "pastures_each": pastures, "pokemon_each": pokemon, "machines_each": machines,
              "heap": HEAP, "report": parse_report(report), "raw_report": report, "tps": tps, "health": health, "log_errors": errors}
    if profiles:
        kept = OUT / f"{name}.sparkprofile"
        shutil.copy2(profiles[-1], kept)
        result["profile"] = summarise(kept)
    (OUT / f"{name}.json").write_text(json.dumps(result, indent=1))
    return result


def write_doc(results, settle, measure):
    lines = ["# Performance", "",
             f"Generated by `tools/perf_test.py`: the island world template, a {HEAP} heap, {settle} s to settle, then {measure} s measured tick by tick with spark profiling the server thread.",
             "Players are the self-tests' fake clients: the server ticks them, loads their chunks and spawns Pokémon around them. One walks a long ring round the island when there are several. They send and encode no packets and fight no battles, so read this as the cost of the world, not of the network.", "",
             "A tick has 50 ms; over that, the server falls behind.", "",
             "| Scenario | Players | Pastures × Pokémon each | Machines each | Mean tick | p95 | p99 | Worst | Ticks over 50 ms | Pokémon pastured / wild | Working machines | Loaded chunks |",
             "| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |"]
    for r in results:
        p = r["report"]
        lines.append(f"| {r['scenario']} | {r['players']} | {r['pastures_each']} × {r['pokemon_each']} | {r['machines_each']} | {p['mean_ms']:.1f} ms | {p['p95_ms']:.1f} ms | {p['p99_ms']:.1f} ms | "
                     f"{p['max_ms']:.0f} ms | {int(p['spikes_over_50ms'])} of {int(p['ticks'])} | {int(p['pastured'])} / {int(p['wild'])} | {int(p['working'])} of {int(p['machines'])} | {int(p['chunks'])} |")
    for r in results:
        if "profile" not in r:
            continue
        prof = r["profile"]
        total = sum(ms for g, ms in prof["groups"].items() if g != "Waiting for the next tick") or 1
        prof["groups"].pop("Waiting for the next tick", None)
        lines += ["", f"## Where the server thread's busy time goes: {r['scenario']}", "", "| Part | Share |", "| --- | --- |"]
        lines += [f"| {g} | {ms / total * 100:.1f} % |" for g, ms in prof["groups"].items()]
        lines += ["", "Companion mod, heaviest methods including what they call:", ""]
        lines += [f"- `{n}`: {ms / total * 100:.2f} %" for n, ms in prof["companion_top"]]
    (ROOT / "docs/performance.md").write_text("\n".join(lines) + "\n")


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--scenarios", nargs="+", default=list(SCENARIOS), choices=list(SCENARIOS))
    parser.add_argument("--settle", type=int, default=120)
    parser.add_argument("--measure", type=int, default=180)
    parser.add_argument("--no-doc", action="store_true")
    args = parser.parse_args()
    if not eula_accepted():
        sys.exit("build/server/eula.txt does not accept the Minecraft EULA yet (see tools/selftest.py --accept-eula).")
    if not (TEMPLATE / "level.dat").exists():
        sys.exit("no build/world-template: run tools/make_world.py first")
    OUT.mkdir(parents=True, exist_ok=True)
    recipe = json.loads((ROOT / "world/world.json").read_text())
    args_file = prepare()
    results = [run(name, SCENARIOS[name], args_file, recipe, args.settle, args.measure) for name in args.scenarios]
    if not args.no_doc:
        write_doc(results, args.settle, args.measure)
    print("PERF TEST DONE")


if __name__ == "__main__":
    main()