"""探针：Sodium 下区块到底有没有进它的 tracker（diagnose sodium tracker）。

用 Python 层的实例 + build.sh 的产物，跑一次 setblocks + debug，读 Sodium 自己的状态。
用法: python3 tools/probe_sodium.py [mods|empty]
"""
import json
import math
import os
import subprocess
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from mcfold.render import SECTION_VOLUME  # noqa: E402

STORE = Path("/home/newqqagent/mcview/store/vanilla-1.20.1")
OUT = STORE / ".mcfold"
INST = Path("/tmp/exp-inst-root/demo-shader")
USE_MODS = (sys.argv[1] if len(sys.argv) > 1 else "mods") == "mods"
MODS = str(INST / "mods") if USE_MODS else "/tmp/empty-mods"
PNG = sys.argv[2] if len(sys.argv) > 2 else "/tmp/sodium-probe.png"


def classpath() -> str:
    libs = [l.strip() for l in (OUT / "classpath.txt").read_text().splitlines()
            if l.strip() and "/org/ow2/asm/" not in l]
    fab = sorted(str(p) for p in (OUT / "fabric").glob("*.jar"))
    return os.pathsep.join(libs + fab + [
        str(STORE / "versions/1.20.1/1.20.1.jar"),
        str(OUT / "classes-intermediary.jar"),
    ])


def rle(v):
    out = []
    for x in v:
        if out and out[-1][0] == x:
            out[-1][1] += 1
        else:
            out.append([x, 1])
    return out


def main() -> int:
    Path("/tmp/empty-mods").mkdir(exist_ok=True)
    proc = subprocess.Popen(
        ["java", "-Xmx2000M", "-Djava.awt.headless=true",
         f"-Djava.library.path={OUT / 'natives'}",
         f"-Dmcfold.names={OUT / 'names-intermediary.json'}",
         f"-Dfabric.modsFolder={MODS}",
         "-cp", classpath(), "LoaderBoot",
         "--gameDir", str(INST),
         "--assets", str(STORE / "assets"), "--assetIndex", "1.20.1"],
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
        text=True, bufsize=1,
        env={**os.environ, "DISPLAY": ":98", "HOME": str(INST)})

    def rpc(obj):
        proc.stdin.write(json.dumps(obj) + "\n")
        proc.stdin.flush()
        while True:
            line = proc.stdout.readline()
            if not line:
                return None
            s = line.strip()
            if s.startswith("{") and f'"id":{obj["id"]}' in s:
                return json.loads(s)

    t0 = time.time()
    while time.time() - t0 < 300:
        line = proc.stdout.readline()
        if not line:
            break
        if '"ready"' in line:
            break

    try:
        cap = rpc({"id": 1, "cmd": "capacity"})
        ids = rpc({"id": 2, "cmd": "resolve", "blocks": [
            "minecraft:stone", "minecraft:grass_block[snowy=false]",
            "minecraft:oak_log[axis=y]", "minecraft:oak_leaves[persistent=true]"]})["ids"]
        stone, grass, log, leaves = ids
        rpc({"id": 3, "cmd": "viewport", "cx": 0, "cz": 0, "radius": 12})
        rpc({"id": 4, "cmd": "view", "hideGui": True})

        secs = []
        for idx in range(cap["sectionsCount"]):
            y0 = idx * 16 - 64
            v = [0] * SECTION_VOLUME
            for yy in range(16):
                for zz in range(16):
                    for xx in range(16):
                        wy = y0 + yy
                        h = 68 + int(2.5 * math.sin(xx * 0.4) * math.cos(zz * 0.4))
                        i = yy * 256 + zz * 16 + xx
                        if wy < h - 4:
                            v[i] = stone
                        elif wy == h - 4:
                            v[i] = grass
                        elif wy == h - 3:
                            if (xx + zz) % 7 == 0:
                                v[i] = log
                            elif (xx * zz) % 5 == 0:
                                v[i] = leaves
            secs.append({"index": idx, "rle": rle(v)})
        n = 10
        for cx in range(-3, 4):
            for cz in range(-3, 4):
                rpc({"id": n, "cmd": "setblocks", "cx": cx, "cz": cz, "sections": secs})
                n += 1

        dbg = rpc({"id": 89, "cmd": "debug"})
        print("  ── debug ──")
        for k, v in (dbg or {}).items():
            if k != "image":
                print(f"    {k} = {v}")

        rpc({"id": 90, "cmd": "camera", "x": 8.0, "y": 72.0, "z": 8.0,
             "yaw": 42.0, "pitch": 30.0})
        res = rpc({"id": 91, "cmd": "render", "frames": 120, "out": PNG})
        print(f"  vanilla计数: chunks={res.get('chunksRendered')} "
              f"colors={res.get('distinctColors')} → {PNG}")
        return 0
    finally:
        try:
            rpc({"id": 999, "cmd": "shutdown"})
            proc.wait(timeout=20)
        except Exception:
            proc.kill()


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