"""Try to solve the checksum/counter scheme for one arbitration ID.

Needed before we can synthesise a message rather than replay it: change a
signal (say road speed in 0x1A0) and the checksum has to be recomputed or
the receiver rejects the frame. files/PTCAN_protocol.md notes that the
one's-complement scheme that works for several IDs only matches 0x1A0
60-75% of the time, so this brute-forces the remaining variants.

Usage:
    python tools/solve_checksum.py captures/replay_car_to_eps_20260829.csv 0x1A0
"""
import argparse
import csv
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "files"))
from decode_ptcan import ocsum  # noqa: E402


def load_frames(path: Path, target: int) -> list[bytes]:
    out = []
    with open(path, newline="") as f:
        for row in csv.DictReader(f):
            if int(row["arbitration_id"], 16) == target:
                out.append(bytes.fromhex(row["data_hex"]))
    return out


def try_scheme(frames, cs_index, fn) -> float:
    ok = 0
    for d in frames:
        if cs_index >= len(d):
            continue
        rest = [d[i] for i in range(len(d)) if i != cs_index]
        if fn(rest, d) == d[cs_index]:
            ok += 1
    return ok / len(frames) if frames else 0.0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("capture")
    ap.add_argument("arbitration_id")
    a = ap.parse_args()

    target = int(a.arbitration_id, 16)
    frames = load_frames(Path(a.capture), target)
    if not frames:
        print(f"No frames for 0x{target:03X}")
        return 1
    print(f"0x{target:03X}: {len(frames)} frames, {len(set(frames))} distinct, dlc={len(frames[0])}\n")

    # Which byte positions actually move? A checksum byte should look random.
    n = len(frames[0])
    print("byte  distinct values  (a checksum looks high-entropy, a counter cycles)")
    for i in range(n):
        vals = Counter(d[i] for d in frames if len(d) > i)
        sample = " ".join(f"{v:02X}" for v, _ in vals.most_common(6))
        print(f"  b{i}   {len(vals):3d}   {sample}")
    print()

    best = []
    for cs in range(n):
        # Plain sums / xors, with and without a per-ID constant.
        for name, base in (
            ("ocsum", lambda r, d: ocsum(r)),
            ("sum", lambda r, d: sum(r) & 0xFF),
            ("xor", lambda r, d: __import__("functools").reduce(lambda x, y: x ^ y, r, 0)),
        ):
            for const in range(256):
                fn = (lambda b, c: (lambda r, d: (b(r, d) + c) & 0xFF))(base, const)
                score = try_scheme(frames, cs, fn)
                if score > 0.97:
                    best.append((score, cs, f"{name} + 0x{const:02X}"))

    if best:
        best.sort(reverse=True)
        print("Schemes matching >97% of frames:")
        for score, cs, desc in best[:10]:
            print(f"  byte {cs} = {desc}   ({score*100:.1f}%)")
    else:
        print("No simple sum/xor scheme fits >97%.")
        print("Best partial matches:")
        partial = []
        for cs in range(n):
            for name, base in (("ocsum", lambda r, d: ocsum(r)), ("sum", lambda r, d: sum(r) & 0xFF)):
                for const in range(256):
                    fn = (lambda b, c: (lambda r, d: (b(r, d) + c) & 0xFF))(base, const)
                    partial.append((try_scheme(frames, cs, fn), cs, f"{name} + 0x{const:02X}"))
        partial.sort(reverse=True)
        for score, cs, desc in partial[:5]:
            print(f"  byte {cs} = {desc}   ({score*100:.1f}%)")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())