"""Name functions in the stripped VREND.MNG by matching relocation-invariant
signatures harvested from the DPL3 VRENDER .O object files.

For each .O .text function we mask (wildcard) every 32-bit word that a COFF
relocation touches -- those hold call targets / data addresses that the linker
rewrites -- and keep the fixed opcode+register skeleton. A function that carried
over into VREND.MNG (even relocated to a new address) still matches its
skeleton. Word granularity throughout (i860 is fixed 4-byte, word-aligned)."""
import sys, struct, os, glob

def parse_coff(path):
    d = open(path, 'rb').read()
    magic, nscns, timdat, symptr, nsyms, opthdr, flags = struct.unpack_from('<HHIIIHH', d, 0)
    secs = []
    base = 20 + opthdr
    for i in range(nscns):
        o = base + i*40
        name = d[o:o+8].split(b'\0')[0].decode('latin1')
        paddr, vaddr, size, scnptr, relptr, lnnoptr, nreloc, nlnno, sflags = \
            struct.unpack_from('<IIIIIIHHI', d, o+8)
        secs.append(dict(name=name, vaddr=vaddr, size=size, scnptr=scnptr,
                         relptr=relptr, nreloc=nreloc))
    # symbols
    strbase = symptr + nsyms*18
    def symname(o):
        z, offs = struct.unpack_from('<II', d, o)
        if z == 0:
            e = d.index(b'\0', strbase+offs); return d[strbase+offs:e].decode('latin1')
        return d[o:o+8].split(b'\0')[0].decode('latin1')
    syms = []; i = 0
    while i < nsyms:
        o = symptr + i*18
        nm = symname(o)
        value, scnum, typ, sclass, naux = struct.unpack_from('<iHHBB', d, o+8)
        syms.append((nm, value, scnum, sclass)); i += 1 + naux
    return d, secs, syms

def text_section(secs):
    for s in secs:
        if s['name'] == '.text': return s
    return None

def sig_for_functions(path, min_words=6):
    """yield (name, [word|None ...]) skeletons for each .text function."""
    d, secs, syms = parse_coff(path)
    t = text_section(secs)
    if not t or t['size'] == 0: return
    text = d[t['scnptr']:t['scnptr']+t['size']]
    nwords = t['size'] // 4
    reloc_word = set()
    for r in range(t['nreloc']):
        rva = struct.unpack_from('<I', d, t['relptr'] + r*10)[0]
        reloc_word.add(rva // 4)                      # wildcard the whole word
    # function boundaries = .text EXT/STAT symbols, sorted
    fsyms = sorted([(v, nm) for (nm, v, sn, sc) in syms
                    if sn == 1 and sc in (2, 3)], key=lambda x: x[0])
    for i, (v, nm) in enumerate(fsyms):
        end = fsyms[i+1][0] if i+1 < len(fsyms) else t['size']
        w0, w1 = v // 4, end // 4
        if w1 - w0 < min_words: continue
        skel = []
        for w in range(w0, w1):
            if w in reloc_word: skel.append(None)
            else: skel.append(struct.unpack_from('<I', text, w*4)[0])
        if sum(1 for x in skel if x is not None) < min_words: continue
        yield nm, skel

def load_mng_text(path):
    d = open(path, 'rb').read()
    tsize, dsize, bsize = struct.unpack_from('<III', d, 0)
    text = d[0x0c:0x0c+tsize]
    return text, tsize

def build_index(text):
    idx = {}
    for p in range(len(text)//4):
        w = struct.unpack_from('<I', text, p*4)[0]
        idx.setdefault(w, []).append(p)
    return idx

def match(skel, text, idx, thresh=0.75):
    """Anchor on the first concrete word, then SCORE the concrete words.
    Returns (best_addr, best_frac, unique) so drifted functions still match on
    their stable prologue/skeleton."""
    k = next((i for i, x in enumerate(skel) if x is not None), None)
    if k is None: return None
    concrete = [i for i, x in enumerate(skel) if x is not None]
    best = (None, 0.0); second = 0.0
    for p in idx.get(skel[k], []):
        start = p - k
        if start < 0 or (start + len(skel))*4 > len(text): continue
        good = 0
        for i in concrete:
            if struct.unpack_from('<I', text, (start+i)*4)[0] == skel[i]:
                good += 1
        frac = good / len(concrete)
        if frac > best[1]:
            second = best[1]; best = (start*4, frac)
        elif frac > second:
            second = frac
    if best[0] is not None and best[1] >= thresh:
        return (best[0], best[1], best[1] > second + 0.05)
    return None

def main():
    mng = sys.argv[1] if len(sys.argv) > 1 else r'C:\VWE\TeslaRel410\sda4\RPLIVE\VREND.MNG'
    odir = r'C:\VWE\TeslaRel410\sda4\DPL3\VRENDER'
    text, tsize = load_mng_text(mng)
    idx = build_index(text)
    print(f"VREND.MNG .text = {tsize:#x} bytes ({tsize//4} words); {os.path.basename(mng)}")
    named = {}   # addr -> (name, frac)
    nsig = 0
    for opath in sorted(glob.glob(os.path.join(odir, '*.O')) +
                        glob.glob(os.path.join(odir, 'PXPL5SUP', '*.O'))):
        for nm, skel in sig_for_functions(opath):
            nsig += 1
            r = match(skel, text, idx)
            if r is None: continue
            addr, frac, uniq = r
            if not uniq: nm = nm + "?"
            if addr not in named or frac > named[addr][1]:
                named[addr] = (nm, frac)
    strong = sum(1 for (n, f) in named.values() if f >= 0.95)
    print(f"signatures tried={nsig}, named={len(named)} "
          f"(exact/near-exact >=0.95: {strong})")
    for addr in sorted(named):
        nm, frac = named[addr]
        print(f"  text+{addr:#08x}  (mem {0x0c+addr:#08x})  {frac:0.2f}  {nm}")

if __name__ == '__main__':
    main()