#!/usr/bin/env python3
"""Backfill authentic TEST_1/ZBUF_1 words into an existing scheduler fixture.

Ch357-Ch404 fixtures used placeholder words because the board wrapper supplied
ZTE/ZMSK constants to the external LPDDR ROP.  Ch405 carries that state from
the feeder list, so accepted fixture geometry can be retained byte-for-byte
while only header words 3 and 4 are restored from the capture.
"""
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.normpath(os.path.join(HERE, ".."))
DATA = os.path.join(ROOT, "sim", "data", "top_psmct32_raster_demo")
sys.path.insert(0, HERE)
import gs_make_sh3_multidraw_fixture as MD


def rows(path):
    with open(path, encoding="utf-8") as f:
        for line in f:
            s = line.strip()
            if s and not s.startswith("#") and not s.startswith("META"):
                fields = s.split()
                if fields[0].isdigit():
                    yield int(fields[1]), fields[8]


def main(argv):
    if len(argv) < 3:
        raise SystemExit("usage: patch_runtime_sched_zstate.py dump.gs.zst epochs.txt [epochs.txt ...]")
    dump = argv[1]
    tables = argv[2:]
    entries = [(idx, name) for table in tables for idx, name in rows(table)]
    got, _ = MD.load_draws(dump, [idx for idx, _ in entries])
    for idx, name in entries:
        if idx not in got:
            raise SystemExit(f"missing captured draw idx{idx}")
        path = os.path.join(DATA, name)
        with open(path, encoding="utf-8") as f:
            lines = f.readlines()
        slots = [i for i, line in enumerate(lines)
                 if line.strip() and not line.lstrip().startswith("//")]
        if len(slots) < 5:
            raise SystemExit(f"{name}: truncated staging image")
        state = got[idx]["state"]
        lines[slots[3]] = f"{state['test'] & 0xffffffffffffffff:016x}\n"
        lines[slots[4]] = f"{state['zbuf'] & 0xffffffffffffffff:016x}\n"
        with open(path, "w", encoding="utf-8") as f:
            f.writelines(lines)
        print(f"[zstate] {name}: idx{idx} TEST={state['test']:016x} ZBUF={state['zbuf']:016x}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))