#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# check_module_layout.py - mechanically compare the frozen CUBB2/CBD2 structs
# in kernel/gsp_cuda_bridge.c (module side) and cuda_bridge.c (daemon side).
#
# Parses both struct bodies, computes x86_64 natural-alignment offsets and
# sizes, and reports the first mismatch. The compile-time static asserts in
# each C file pin their own layout; this check pins the two sides to each
# other. Exit code 0 = identical.

import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
MODULE = os.path.normpath(os.path.join(HERE, "..", "..", "kernel",
                                       "gsp_cuda_bridge.c"))
DAEMON = os.path.join(HERE, "cuda_bridge.c")

TYPE_SIZE = {
    "u8": 1, "u32": 4, "u64": 8,
    "uint8_t": 1, "uint32_t": 4, "uint64_t": 8,
    "char": 1,
}
CONST = {"CUBB2_NAME_MAX": 64}


def parse_struct(path, name):
    text = open(path).read()
    m = re.search(r"struct\s+%s\s*\{([^{}]*)\}\s*;" % re.escape(name), text,
                  re.S)
    if not m:
        m = re.search(r"typedef\s+struct\s*\{([^{}]*)\}\s*%s_t\s*;"
                      % re.escape(name), text, re.S)
    if not m:
        raise SystemExit("%s: struct %s not found" % (path, name))
    fields = []
    for raw in m.group(1).splitlines():
        line = raw.split("/*")[0].strip()
        if not line or line.startswith("*"):
            continue
        m2 = re.match(r"^([A-Za-z_][A-Za-z0-9_]*)\s+([A-Za-z_][A-Za-z0-9_]*)"
                      r"\s*(?:\[\s*(.+?)\s*\])?\s*;$", line)
        if not m2:
            raise SystemExit("%s: cannot parse field %r" % (path, line))
        typ, fname, count_expr = m2.group(1), m2.group(2), m2.group(3)
        count = 1
        if count_expr is not None:
            count = int(eval(count_expr, {"__builtins__": {}}, CONST))
        if typ not in TYPE_SIZE:
            raise SystemExit("%s: unknown type %s" % (path, typ))
        fields.append((typ, fname, count))
    return fields


def layout(fields):
    out = []
    off = 0
    align = 1
    for typ, fname, count in fields:
        sz = TYPE_SIZE[typ]
        al = sz
        off = (off + al - 1) // al * al
        out.append((fname, typ, count, off, sz * count))
        off += sz * count
        align = max(align, al)
    total = (off + align - 1) // align * align
    return out, total


def compare(name, mpath, dpath):
    mod = layout(parse_struct(mpath, name))
    dae = layout(parse_struct(dpath, name))
    mfields, mtotal = mod
    dfields, dtotal = dae
    ok = True
    print("=== struct %s (module %s | daemon %s) ===" % (name, mpath, dpath))
    if mtotal != dtotal:
        print("TOTAL MISMATCH: module=%d daemon=%d" % (mtotal, dtotal))
        ok = False
    if len(mfields) != len(dfields):
        print("FIELD COUNT MISMATCH: module=%d daemon=%d"
              % (len(mfields), len(dfields)))
        return False
    print("%-20s %-10s %-18s %-18s" % ("field", "type", "module off/size",
                                       "daemon off/size"))
    for (mn, mt, mc, mo, ms), (dn, dt, dc, do, ds) in zip(mfields, dfields):
        same = (mn == dn and mc == dc and mo == do and ms == ds)
        print("%-20s %-10s 0x%02x/%-12d 0x%02x/%-12d%s"
              % (mn, "%s|%s" % (mt, dt), mo, ms, do, ds,
                 "" if same else "  <-- MISMATCH"))
        if not same:
            ok = False
    print("total: module=0x%x daemon=0x%x %s"
          % (mtotal, dtotal, "MATCH" if mtotal == dtotal else "MISMATCH"))
    return ok


def main():
    ok = True
    for name in ("cubb2_job", "cubb2_done"):
        ok = compare(name, MODULE, DAEMON) and ok
    print("LAYOUT-CHECK %s" % ("PASS" if ok else "FAIL"))
    return 0 if ok else 1


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