#!/usr/bin/env python3
"""Phase D checksum cracker: 16-bit/ones-complement/chained/DJB2/Murmur batch.

Covers what Phases A-C did not: every candidate must match ALL 3 builds.
Stdlib only, read-only on stock .bin files.

  Quick smoke:  python3 tools/lofi_checksum_phaseD.py --limit 30
  Full run:     python3 tools/lofi_checksum_phaseD.py --jobs $(nproc) --out /tmp/opencode/ck-phaseD.txt

Families:
  CRC16: ARC/Modbus/CCITT-FALSE/XMODEM/Kermit/DNP/UMTS/BINHEX + halved-file pairs
  Ones-complement: 16-bit word sums (LE/BE, folded, complemented)
  Word sums: 16-bit LE/BE, BSD rotate, SysV folded
  Chained per-sect: crc32-of-crcs, sum-of-crcs, xor-of-crcs
  Odd hashes: FNV-1 (not 1a), DJB2, Murmur3-x86-32
"""
import argparse, binascii, multiprocessing as mp
import os, struct, sys, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from lofi_checksum_crack import load_images, span_bufs, SPANS, _FIELDMODE_IDX

CRC16_PARAMS = [  # (name, poly_trunc, init, refin, xorout, check_of_"123456789")
    # NOTE: refin polys are given in bit-reversed (truncated) form, used as-is.
    ('ARC',       0xA001, 0x0000, True,  0x0000, 0xBB3D),
    ('MODBUS',    0xA001, 0xFFFF, True,  0x0000, 0x4B37),
    ('USB',       0xA001, 0xFFFF, True,  0xFFFF, 0xB4C8),
    ('CCITT-F',   0x1021, 0xFFFF, False, 0x0000, 0x29B1),
    ('XMODEM',    0x1021, 0x0000, False, 0x0000, 0x31C3),
    ('KERMIT',    0x8408, 0x0000, True,  0x0000, 0x2189),
    ('X25',       0x8408, 0xFFFF, True,  0xFFFF, 0x906E),
    ('DNP',       0xA6BC, 0x0000, True,  0xFFFF, 0xEA82),
    ('GENIBUS',   0x1021, 0xFFFF, False, 0xFFFF, 0xD64E),
    ('GSM',       0x1021, 0x0000, False, 0xFFFF, 0xCE3C),
]

_tables16 = {}
def _tab16(poly_trunc, refin):
    key = (poly_trunc, refin)
    t = _tables16.get(key)
    if t is not None:
        return t
    if refin:
        # poly_trunc is already bit-reversed (e.g. 0xA001); use as-is.
        p = poly_trunc
        t = []
        for i in range(256):
            c = i
            for _ in range(8):
                c = (c >> 1) ^ p if c & 1 else c >> 1
            t.append(c & 0xFFFF)
    else:
        p = poly_trunc
        t = []
        for i in range(256):
            c = i << 8
            for _ in range(8):
                c = ((c << 1) ^ p) & 0xFFFF if c & 0x8000 else (c << 1) & 0xFFFF
            t.append(c)
    _tables16[key] = t
    return t

def crc16(buf, poly, init, refin, xorout):
    tab = _tab16(poly, refin)
    crc = init
    if refin:
        for b in buf:
            crc = tab[(crc ^ b) & 0xFF] ^ (crc >> 8)
    else:
        for b in buf:
            crc = tab[((crc >> 8) ^ b) & 0xFF] ^ ((crc << 8) & 0xFFFF)
    return (crc ^ xorout) & 0xFFFF

import array

def _even(buf):
    return buf if len(buf) % 2 == 0 else buf + b'\x00'

def _words16(buf, endian):
    a = array.array('H', _even(buf))
    if endian == 'big':
        a.byteswap()
    return a

def ones_complement16(buf, endian):
    mv = _words16(buf, endian)
    s = sum(mv) & 0xFFFFFFFFFFFFFFFF
    while s >> 16:
        s = (s & 0xFFFF) + (s >> 16)
    return (~s) & 0xFFFF

def wordsum16(buf, endian):
    return sum(_words16(buf, endian)) & 0xFFFF

def bsd_sum(buf):
    s = 0
    for b in buf:
        s = ((s >> 1) | ((s & 1) << 15)) & 0xFFFF
        s = (s + b) & 0xFFFF
    return s

def sysv_sum(buf):
    s = sum(buf)
    s = (s & 0xFFFF) + ((s >> 16) & 0xFFFF)
    s = (s & 0xFFFF) + ((s >> 16) & 0xFFFF)
    return s & 0xFFFF

def fnv1_32(buf):
    h = 0x811C9DC5
    for b in buf:
        h = (h * 0x01000193) & 0xFFFFFFFF
        h ^= b
    return h

def djb2(buf):
    h = 5381
    for b in buf:
        h = ((h * 33) + b) & 0xFFFFFFFF
    return h

def murmur3_x86_32(buf, seed=0):
    h = seed
    n = len(buf) & ~3
    for i in range(0, n, 4):
        k = struct.unpack('<I', buf[i:i+4])[0]
        k = (k * 0xCC9E2D51) & 0xFFFFFFFF
        k = ((k << 15) | (k >> 17)) & 0xFFFFFFFF
        k = (k * 0x1B873593) & 0xFFFFFFFF
        h ^= k
        h = ((h << 13) | (h >> 19)) & 0xFFFFFFFF
        h = (h * 5 + 0xE6546B64) & 0xFFFFFFFF
    tail = buf[n:]
    k = 0
    for i, b in enumerate(tail):
        k |= b << (8 * i)
    if tail:
        k = (k * 0xCC9E2D51) & 0xFFFFFFFF
        k = ((k << 15) | (k >> 17)) & 0xFFFFFFFF
        k = (k * 0x1B873593) & 0xFFFFFFFF
        h ^= k
    h ^= len(buf)
    h ^= h >> 16
    h = (h * 0x85EBCA6B) & 0xFFFFFFFF
    h ^= h >> 13
    h = (h * 0xC2B2AE35) & 0xFFFFFFFF
    h ^= h >> 16
    return h

def match16(stored, v):
    """How a 16-bit value could sit in the 32-bit field."""
    lo, hi = stored & 0xFFFF, (stored >> 16) & 0xFFFF
    if lo == v == hi:
        return 'both16'
    if lo == v:
        return 'lo16'
    if hi == v:
        return 'hi16'
    if stored == v:
        return 'full32eq16'
    if stored == ((v << 16) | v):
        return 'duplicated16'
    return None

def build_jobs():
    jobs = []
    for name, poly, init, refin, xorout, _check in CRC16_PARAMS:
        for span in SPANS:
            for fmode in ('asis', 'zeroed', 'ff'):
                jobs.append(('crc16', name, (poly, init, refin, xorout), span, fmode))
    for span in SPANS:
        for fmode in ('asis', 'zeroed', 'ff'):
            for nm in ('ones_le', 'ones_be', 'wsum16_le', 'wsum16_be',
                       'bsd', 'sysv', 'fnv1', 'djb2', 'murmur0', 'murmurF'):
                jobs.append(('fast32', nm, (), span, fmode))
    for span in SPANS:
        for fmode in ('asis', 'zeroed'):
            for nm in ('chain_crc_of_crcs', 'chain_sum_of_crcs', 'chain_xor_of_crcs',
                       'halved_crc16_pair'):
                jobs.append(('chain', nm, (), span, fmode))
    return jobs

_G = None
def _init(blobs):
    global _G
    _G = blobs  # [(ver, span->bufs[3], ck)]

FAST32 = {
    'fnv1': fnv1_32, 'djb2': djb2,
    'murmur0': lambda b: murmur3_x86_32(b, 0),
    'murmurF': lambda b: murmur3_x86_32(b, 0xFFFFFFFF),
}

def _run(job):
    kind, name, params, span, fmode = job
    idx = _FIELDMODE_IDX[fmode]
    bufs = [(span_bufs_alias(d, span, idx), ck) for _, d, ck in _G]
    if kind == 'crc16':
        poly, init, refin, xorout = params
        got = [crc16(b, poly, init, refin, xorout) for b, _ in bufs]
        how = [match16(ck, v) for (_, ck), v in zip(bufs, got)]
        if all(how):
            return f'HIT crc16/{name} span={span} field={fmode} val={got[0]:04X} as={how[0]}'
        masks = [(ck ^ v) & 0xFFFF for (_, ck), v in zip(bufs, got)]
        if masks[0] == masks[1] == masks[2]:
            return (f'MASK16 crc16/{name} span={span} field={fmode} '
                    f'mask={masks[0]:04X}')
        return None
    if kind == 'fast32':
        if name == 'ones_le':
            got = [ones_complement16(b, 'little') for b, _ in bufs]
            how = [match16(ck, v) for (_, ck), v in zip(bufs, got)]
            if all(how):
                return f'HIT ones-complement-LE span={span} field={fmode} as={how[0]}'
            return None
        if name == 'ones_be':
            got = [ones_complement16(b, 'big') for b, _ in bufs]
            how = [match16(ck, v) for (_, ck), v in zip(bufs, got)]
            if all(how):
                return f'HIT ones-complement-BE span={span} field={fmode} as={how[0]}'
            return None
        if name == 'wsum16_le':
            got = [wordsum16(b, 'little') for b, _ in bufs]
        elif name == 'wsum16_be':
            got = [wordsum16(b, 'big') for b, _ in bufs]
        elif name == 'bsd':
            got = [bsd_sum(b) for b, _ in bufs]
        elif name == 'sysv':
            got = [sysv_sum(b) for b, _ in bufs]
        else:
            fn = FAST32[name]
            got = [fn(b) for b, _ in bufs]
            if all(ck == v for (_, ck), v in zip(bufs, got)):
                return f'HIT {name} span={span} field={fmode} val={got[0]:08X}'
            masks = [(ck ^ v) & 0xFFFFFFFF for (_, ck), v in zip(bufs, got)]
            if masks[0] == masks[1] == masks[2]:
                return f'MASK32 {name} span={span} field={fmode} mask={masks[0]:08X}'
            return None
        how = [match16(ck, v) for (_, ck), v in zip(bufs, got)]
        if all(how):
            return f'HIT {name} span={span} field={fmode} val={got[0]:04X} as={how[0]}'
        return None
    # chained per-sect schemes
    nsects = [struct.unpack('<I', d[0x50:0x54])[0] for _, d, _ in _G]
    crcs_per_image = []
    for (_, d, _), n in zip(_G, nsects):
        off, lst = 0x54, []
        for _ in range(n):
            ln = struct.unpack('<I', d[off+8:off+12])[0]
            lst.append(binascii.crc32(d[off+12:off+12+ln]) & 0xFFFFFFFF)
            off += 12 + ln
        crcs_per_image.append(lst)
    if name == 'chain_crc_of_crcs':
        got = [binascii.crc32(struct.pack(f'<{len(c)}I', *c)) & 0xFFFFFFFF
               for c in crcs_per_image]
    elif name == 'chain_sum_of_crcs':
        got = [sum(c) & 0xFFFFFFFF for c in crcs_per_image]
    elif name == 'chain_xor_of_crcs':
        got = []
        for c in crcs_per_image:
            x = 0
            for v in c:
                x ^= v
            got.append(x)
    elif name == 'halved_crc16_pair':
        got = []
        for b, _ in bufs:
            h = len(b) // 2
            a = crc16(b[:h], 0xA001, 0, True, 0)
            c = crc16(b[h:], 0xA001, 0, True, 0)
            got.append(((a << 16) | c) & 0xFFFFFFFF)
    else:
        return None
    if all(ck == v for (_, ck), v in zip(bufs, got)):
        return f'HIT {name} span={span} field={fmode} val={got[0]:08X}'
    masks = [(ck ^ v) & 0xFFFFFFFF for (_, ck), v in zip(bufs, got)]
    if masks[0] == masks[1] == masks[2]:
        return f'MASK32 {name} span={span} field={fmode} mask={masks[0]:08X}'
    return None

def span_bufs_alias(d, span, idx):
    return span_bufs(d)[span][idx]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--jobs', type=int, default=max(1, (os.cpu_count() or 4) - 1))
    ap.add_argument('--out', default='/tmp/opencode/ck-phaseD.txt')
    ap.add_argument('--limit', type=int, default=None)
    a = ap.parse_args()
    blobs = load_images()
    packed = [(ver, d, ck) for ver, d, ck in blobs]
    for ver, d, ck in blobs:
        print(f'{ver}: size={len(d)} stored_ck={ck:08X}', flush=True)
    jobs = build_jobs()
    if a.limit:
        jobs = jobs[:a.limit]
    print(f'Phase D: {len(jobs)} combos, jobs={a.jobs}', flush=True)
    t0 = time.time()
    found, done = [], 0
    total = len(jobs)
    with mp.Pool(a.jobs, initializer=_init, initargs=(packed,)) as pool:
        for res in pool.imap_unordered(_run, jobs, chunksize=4):
            done += 1
            if res:
                print(f'  *** {res}', flush=True)
                found.append(res)
                with open(a.out, 'a') as fh:
                    fh.write(res + '\n')
            if done % 50 == 0 or done == total:
                print(f'  [{done}/{total}] elapsed={time.time()-t0:.0f}s',
                      flush=True)
    print(f'Phase D done: {total} combos in {time.time()-t0:.0f}s, '
          f'{len(found)} candidates.', flush=True)

if __name__ == '__main__':
    main()