import os, sys
from itertools import product
from multiprocessing import Pool, cpu_count
from binascii import unhexlify
import tqdm

def load_samples(folder):
    msgs, crcs = [], []
    for fn in os.listdir(folder):
        with open(os.path.join(folder,fn),'r') as f:
            for L in f:
                if '=' not in L: continue
                h,c = L.strip().split('=',1)
                try:
                    msgs.append(unhexlify(h))
                    crcs.append(int(c,16))
                except:
                    pass
    return msgs, crcs

def calc_crc8(data, poly, init, xorout, ri, ro):
    crc = init
    for b in data:
        if ri: b = int('{:08b}'.format(b)[::-1],2)
        crc ^= b
        for _ in range(8):
            crc = ((crc << 1) ^ poly) & 0xFF if crc & 0x80 else (crc << 1) & 0xFF
    if ro: crc = int('{:08b}'.format(crc)[::-1],2)
    return crc ^ xorout

def test_config(args):
    poly, init, xorout, ri, ro, msgs, crcs = args
    ok = 0
    total = len(msgs)
    for d, chk in zip(msgs, crcs):
        if calc_crc8(d, poly, init, xorout, ri, ro) == chk:
            ok += 1
    return (poly, init, xorout, ri, ro, ok, total)

def main(folder):
    msgs, crcs = load_samples(folder)
    if len(msgs) < 2:
        sys.exit("Need at least 2 samples")
    tasks = [
        (poly, init, xorout, ri, ro, msgs, crcs)
        for poly in range(256) if (poly & 1 or poly == 1)
        for init, xorout, ri, ro in product(range(256), range(256), [False, True], [False, True])
    ]

    print(f"Wasting ressources on {len(tasks):,} configs using {cpu_count()} cores... yea")
    best = []
    target = len(msgs)
    with Pool() as p:
        for poly, init, xorout, ri, ro, ok, total in tqdm.tqdm(p.imap_unordered(test_config, tasks), total=len(tasks)):
            if ok == total:
                print("FOUND FULL FUCKING MATCH ✅", f"poly=0x{poly:02X}", f"init=0x{init:02X}",
                      f"xorout=0x{xorout:02X}", f"ri={ri}", f"ro={ro}")
                return
            # keep top-5 partials
            best.append((ok, poly, init, xorout, ri, ro))
            best = sorted(best, reverse=True)[:5]

    print("\n=== TOP PARTIAL MATCHES ===")
    for ok, poly, init, xorout, ri, ro in best:
        print(f"{ok}/{total} matches – poly=0x{poly:02X}, init=0x{init:02X}, xorout=0x{xorout:02X}, ri={ri}, ro={ro}")
    print("👎 No full match found.")

if __name__ == '__main__':
    if len(sys.argv) != 2:
        sys.exit(f"Usage: python {sys.argv[0]} <samples_folder>")
    main(sys.argv[1])