#!/usr/bin/env python3
"""Same-length UI string swap. Keeps every address stable (Level-0 mod).

Example: lofi_patch_string.py "Lofi-12 XT.bin" out.bin Threshold ThresholX
"""
import argparse, struct, sys, os
sys.path.insert(0, os.path.dirname(__file__))
from lofi_image import parse_image, build_image, find_sect_by_role

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('image')
    ap.add_argument('output')
    ap.add_argument('old')
    ap.add_argument('new')
    ap.add_argument('--sect', type=int, default=None)
    ap.add_argument('--first-only', action='store_true')
    a = ap.parse_args()
    old_b, new_b = a.old.encode(), a.new.encode()
    assert len(old_b) == len(new_b), \
        f"lengths differ ({len(old_b)} vs {len(new_b)}); use equal length to keep addresses stable"
    d = open(a.image, 'rb').read()
    info = parse_image(d)
    tgt = None
    if a.sect is not None:
        tgt = info['sects'][a.sect]
    else:
        tgt = find_sect_by_role(info, 'rodata')
    hits = []
    p = tgt['payload']
    start = 0
    while True:
        i = p.find(old_b, start)
        if i < 0:
            break
        hits.append(i)
        start = i + 1
    if not hits:
        sys.exit(f"'{a.old}' not found in sect{tgt['index']} (addr {tgt['addr']:08X})")
    print(f"found {len(hits)} hit(s) in sect{tgt['index']} at offsets: "
          + ', '.join(f"+{h:#x} (file {tgt['off']+12+h:#x})" for h in hits))
    if a.first_only:
        hits = hits[:1]
    np = bytearray(p)
    for h in hits:
        np[h:h+len(old_b)] = new_b
    payloads = [(s['addr'], bytes(np) if s['index'] == tgt['index'] else s['payload'])
                for s in info['sects']]
    out = build_image(info, payloads)  # recomputes valid checksum
    open(a.output, 'wb').write(out)
    print(f"wrote {a.output} with fresh checksum "
          f"{struct.unpack('<I', out[8:12])[0]:08X} — verify on device.")

if __name__ == '__main__':
    main()