#!/usr/bin/env python3
"""Concatenate runtime-CLUT scheduler fixtures without regenerating accepted epochs.

Usage: merge_runtime_sched.py [--reference-inputs] OUTPUT_TAG INPUT_TAG[@START:END] [...]

All inputs must use the same framebuffer and fixed 512x512 cache geometry.  The
tool renumbers epochs and emits the descriptor/parameter headers consumed by
the shared RTL testbench and host.  By default it copies generated assets under
OUTPUT_TAG.  With --reference-inputs, the host descriptor retains the original
asset names (so repeated resident textures are staged only once), while cheap
local hard-link aliases satisfy the shared RTL testbench's epoch-local naming
contract.
"""
import os
import re
import shutil
import sys

ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))
DATA = os.path.join(ROOT, "sim", "data", "top_psmct32_raster_demo")


def parse_table(tag):
    path = os.path.join(DATA, f"sh3_{tag}_epochs.txt")
    meta = None
    rows = []
    with open(path) as f:
        for line in f:
            if line.startswith("META "):
                meta = line.strip()
            elif line and line[0].isdigit():
                fields = line.split()
                if len(fields) != 14:
                    raise ValueError(f"{path}: expected 14 fields, got {len(fields)}")
                rows.append(fields)
    if meta is None:
        raise ValueError(f"{path}: no META row")
    return meta, rows


def parse_spec(spec):
    if "@" not in spec:
        tag=spec; bounds=None
    else:
        tag,bounds=spec.rsplit("@",1)
    meta,rows=parse_table(tag)
    if bounds is not None:
        a,b=bounds.split(":",1)
        start=int(a) if a else 0; end=int(b) if b else len(rows)
        if start<0 or end<start or end>len(rows):
            raise ValueError(f"{spec}: invalid row slice for {len(rows)} rows")
        rows=rows[start:end]
    return tag,meta,rows


def meta_dict(line):
    p = line.split()[1:]
    return dict(zip(p[0::2], p[1::2]))


def link_or_copy(src, dst):
    """Create a local testbench alias without duplicating large texture files."""
    try:
        os.unlink(dst)
    except FileNotFoundError:
        pass
    try:
        os.link(src, dst)
    except OSError:
        shutil.copyfile(src, dst)


def write_zero_pal(dst):
    """Satisfy the shared RTL harness for direct-color epochs.

    Host descriptors keep '-' so no runtime CLUT upload occurs.  The generic
    testbench still opens an epoch-local palette filename before TEX0.PSM is
    known; a deterministic unused zero table avoids a misleading readmem
    warning without changing the fixture contract.
    """
    with open(dst, "w") as f:
        f.write("// unused direct-color palette placeholder\n")
        for _ in range(256):
            f.write("00000000\n")


def main(argv):
    reference_inputs = "--reference-inputs" in argv
    args = [x for x in argv[1:] if x != "--reference-inputs"]
    if len(args) < 3:
        print(__doc__.strip())
        return 2
    out, tags = args[0], args[1:]
    parsed = [(spec, *parse_spec(spec)) for spec in tags]
    metas = [meta_dict(x[2]) for x in parsed]
    fixed = ("fbpxw", "fbh", "fbwords", "lpddr_tex", "tex_words", "n_beats")
    for key in fixed:
        vals = {m[key] for m in metas}
        if len(vals) != 1:
            raise ValueError(f"input fixtures disagree on {key}: {sorted(vals)}")

    out_rows = []
    nk = 0
    for spec, tag, _meta, rows in parsed:
        for old in rows:
            oldk = int(old[0])
            tex_src = os.path.join(DATA, old[4])
            list_src = os.path.join(DATA, old[8])
            has_pal = old[12] != "-"
            pal_src = os.path.join(DATA, old[12]) if has_pal else None
            tex_alias = f"sh3_{out}{nk}_tex_lpddr.mem"
            list_alias = f"feeder_sh3_{out}{nk}.mem"
            pal_alias = f"sh3_{out}{nk}_pal.mem"
            if reference_inputs:
                tex_name, list_name, pal_name = old[4], old[8], old[12]
                assets = [(tex_src, tex_alias), (list_src, list_alias)]
                if has_pal:
                    assets.append((pal_src, pal_alias))
                for src, name in assets:
                    link_or_copy(src, os.path.join(DATA, name))
                if not has_pal:
                    write_zero_pal(os.path.join(DATA, pal_alias))
            else:
                tex_name, list_name = tex_alias, list_alias
                pal_name = pal_alias if has_pal else "-"
                assets = [(tex_src, tex_name), (list_src, list_name)]
                if has_pal:
                    assets.append((pal_src, pal_name))
                for src, name in assets:
                    shutil.copyfile(src, os.path.join(DATA, name))
                if not has_pal:
                    write_zero_pal(os.path.join(DATA, pal_alias))
            row = list(old)
            row[0] = str(nk)
            row[4], row[8], row[12] = tex_name, list_name, pal_name
            # A sliced input can begin with reuse=1 even though the preceding
            # output row came from another fixture and holds another texture.
            # Never carry that stale residency promise across a merge seam.
            if row[11] == "1" and (not out_rows or out_rows[-1][7] != row[7]):
                row[11] = "0"
            out_rows.append(row)
            nk += 1

    m = metas[0]
    with open(os.path.join(DATA, f"sh3_{out}_epochs.txt"), "w") as f:
        f.write(f"# merged runtime-CLUT scheduler fixture: {tags}\n")
        f.write("# k idx tbp cbp_reloc tex_file lpddr size crc list_file words records reuse pal_file pal_sum32\n")
        f.write("META n_epochs %d fbpxw %s fbh %s fbwords %s lpddr_tex %s tex_words %s n_beats %s\n" %
                (nk, m["fbpxw"], m["fbh"], m["fbwords"], m["lpddr_tex"], m["tex_words"], m["n_beats"]))
        for row in out_rows:
            f.write(" ".join(row) + "\n")

    first_tag=parsed[0][1]
    base_params = open(os.path.join(DATA, f"sh3_{first_tag}_params.vh")).read()
    globals_only = re.split(r"localparam int\s+EP0_NTRIS", base_params, maxsplit=1)[0]
    globals_only = re.sub(r"(localparam int N_EPOCHS\s*=\s*)\d+", rf"\g<1>{nk}", globals_only)
    with open(os.path.join(DATA, f"sh3_{out}_params.vh"), "w") as f:
        f.write(f"// merged runtime scheduler params: {tags}\n")
        f.write(globals_only)
        for k, row in enumerate(out_rows):
            f.write(f"localparam int      EP{k}_NTRIS = {row[10]};\n")
            f.write(f"localparam int      EP{k}_CBP   = {row[3]};\n")
            f.write(f"localparam [31:0]   EP{k}_CRC   = 32'h{int(row[7], 16):08x};\n")
            f.write(f"localparam bit      EP{k}_REUSE = 1'b{row[11]};\n")
    with open(os.path.join(DATA, f"sh3_{out}_epoch_table.vh"), "w") as f:
        f.write(f"// merged runtime scheduler table: {tags}\n")
        for k in range(nk):
            f.write(f"EP_CRC[{k}] = EP{k}_CRC; EP_REC[{k}] = EP{k}_NTRIS; EP_REUSE[{k}] = EP{k}_REUSE;\n")

    for stem in ("bios", "payload"):
        shutil.copyfile(os.path.join(DATA, f"{stem}_sh3_{first_tag}.mem"),
                        os.path.join(DATA, f"{stem}_sh3_{out}.mem"))
    print(f"[merge_sched] {tags} -> {out}: {nk} epochs ({'referenced inputs' if reference_inputs else 'copied assets'})")
    return 0


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