"""Flip SAI/O2 device-enable flags in a TuneECU map, in place.

Byte-boolean array in the flat ROM: 0x53801 = SAI, 0x53818 / 0x53819 = the two
O2 sensors (1 = enabled, 0 = disabled). Found by diffing stock 20188 against the
community "NO SAI, NO O2 SENSORS" reference map -- see DEVICES.md for the full
derivation and confidence levels (SAI ~90%, O2 ~80%).

This edits the *download-format* map (decode -> flip bytes in the flat ROM ->
repack into the decoded map's own layout -> re-encode). It does NOT touch the
ECU flash checksum (out of scope, unsolved -- see the "Editing / export" section
of DEVICES.md), so the output is for the viewer / further analysis only. It is
NOT known to be flashable as-is.

Usage:
    python3 toggle_devices.py 20262Map.hex out/20262-noSAI-noO2.hex
    python3 toggle_devices.py --only-sai 20262Map.hex out/20262-noSAI.hex
"""

from __future__ import annotations

import argparse

from decode_map import decode, encode

SAI = 0x53801
O2_1 = 0x53818
O2_2 = 0x53819

FLAT_MARKER = 0x6F66

_LABELS = {SAI: "SAI", O2_1: "O2 sensor 1", O2_2: "O2 sensor 2"}


def _le(buf: bytes, o: int, n: int) -> int:
    return int.from_bytes(buf[o : o + n], "little")


def unpack(decoded: bytes) -> tuple[bytearray, list[tuple[int, int, int]]]:
    """Reconstruct the flat ROM from a decoded map; also return the (packed_pos,
    dest_offset, length) triples needed to repack it afterwards."""
    i21 = _le(decoded, 28, 2)
    if _le(decoded, i21 + 31, 2) != FLAT_MARKER:
        raise ValueError(f"bad unpack marker at 0x{i21 + 31:X}")
    count = decoded[i21 + 33]
    entries = [
        (_le(decoded, i21 + 34 + k * 8, 4), _le(decoded, i21 + 38 + k * 8, 4))
        for k in range(count)
    ]
    p = i21 + 34 + count * 8
    rom = bytearray(b"\xff" * max(off + ln for off, ln in entries))
    positions = []
    for off, ln in entries:
        rom[off : off + ln] = decoded[p : p + ln]
        positions.append((p, off, ln))
        p += ln
    return rom, positions


def repack(decoded: bytes, rom: bytes, positions: list[tuple[int, int, int]]) -> bytes:
    """Inverse of unpack(): copy the (possibly modified) flat ROM back into the
    decoded map's packed layout."""
    out = bytearray(decoded)
    for p, off, ln in positions:
        out[p : p + ln] = rom[off : off + ln]
    return bytes(out)


def toggle(raw: bytes, addresses: tuple[int, ...]) -> tuple[bytes, list[tuple[int, str, int, int]]]:
    """Return (re-encoded .hex bytes, [(address, label, before, after), ...])."""
    decoded = decode(raw)
    rom, positions = unpack(decoded)
    changes = []
    for addr in addresses:
        before = rom[addr]
        rom[addr] = 0
        changes.append((addr, _LABELS.get(addr, f"0x{addr:X}"), before, 0))
    repacked = repack(decoded, bytes(rom), positions)
    out = encode(repacked)
    assert decode(out) == repacked, "round-trip failed after repack"
    return out, changes


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("infile")
    ap.add_argument("outfile")
    ap.add_argument("--only-sai", action="store_true", help="disable SAI only")
    ap.add_argument("--only-o2", action="store_true", help="disable both O2 sensors only")
    args = ap.parse_args()

    if args.only_sai:
        addrs = (SAI,)
    elif args.only_o2:
        addrs = (O2_1, O2_2)
    else:
        addrs = (SAI, O2_1, O2_2)

    raw = open(args.infile, "rb").read()
    out, changes = toggle(raw, addrs)
    open(args.outfile, "wb").write(out)

    for addr, label, before, after in changes:
        print(f"  0x{addr:05X}  {label:<14} {before} -> {after}")
    print(f"wrote {args.outfile} ({len(out)} bytes)")
    print("NOTE: flash checksum not patched -- not known to be flashable as-is.")
    return 0


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