#!/usr/bin/env python3
"""Create or apply IPS patches (own code, no dependencies).

  python3 tools/ips.py make  ORIGINAL.nes MODIFIED.nes OUT.ips
  python3 tools/ips.py apply ORIGINAL.nes PATCH.ips  OUT.nes
"""
import sys, hashlib

def md5(b): return hashlib.md5(b).hexdigest()

def make(orig, new, out):
    a, b = open(orig, "rb").read(), open(new, "rb").read()
    assert len(b) >= len(a), "modified ROM is shorter than the original (IPS cannot truncate)"
    n = len(a)
    diff = lambda k: k >= n or a[k] != b[k]  # every appended byte is written, so patchers need no zero fill
    recs, i = [], 0
    while i < len(b):
        if not diff(i): i += 1; continue
        j = i
        while j < len(b) and j - i < 0xFFFF and (diff(j) or (j + 1 < len(b) and diff(j + 1))):
            j += 1
        recs.append((i, b[i:j])); i = j
    assert len(b) <= 0x1000000, "file too large for IPS offsets"
    with open(out, "wb") as f:
        f.write(b"PATCH")
        for off, data in recs:
            assert off != 0x454F46, "offset collides with EOF marker"
            f.write(off.to_bytes(3, "big") + len(data).to_bytes(2, "big") + data)
        f.write(b"EOF")
    return a, b, recs

def apply(orig, patch, out):
    rom = bytearray(open(orig, "rb").read()); p = open(patch, "rb").read()
    assert p[:5] == b"PATCH", "not an IPS file"
    i = 5
    while p[i:i + 3] != b"EOF":
        off = int.from_bytes(p[i:i + 3], "big"); size = int.from_bytes(p[i + 3:i + 5], "big"); i += 5
        if size == 0:  # RLE record
            rle = int.from_bytes(p[i:i + 2], "big"); val = p[i + 2]; i += 3
            data = bytes([val]) * rle
        else:
            data = p[i:i + size]; i += size
        if off + len(data) > len(rom): rom.extend(b"\0" * (off + len(data) - len(rom)))
        rom[off:off + len(data)] = data
    open(out, "wb").write(rom)
    return bytes(rom)

if __name__ == "__main__":
    if sys.argv[1] == "make":
        a, b, recs = make(*sys.argv[2:5])
        print(f"original md5 {md5(a)}\nmodified md5 {md5(b)}\n{len(recs)} records:")
        for off, data in recs: print(f"  file offset 0x{off:06X} .. 0x{off + len(data) - 1:06X}  ({len(data)} bytes)")
    elif sys.argv[1] == "apply":
        r = apply(*sys.argv[2:5]); print(f"wrote {sys.argv[4]}  md5 {md5(r)}")