#!/usr/bin/env python3
"""Apply a Pokémon ROM-hack patch to a base ROM (IPS / BPS / UPS).

Format is auto-detected by magic bytes. BPS embeds source+target CRC32, so a
wrong base ROM (or a buggy apply) fails loudly here rather than producing a
silently-corrupt ROM — that is the correctness guarantee for the whole pipeline.

Usage:  romhack-apply.py <patch> <base_rom> <out_rom>
Exit:   0 ok; 2 bad/mismatched input; 3 unsupported format.
xdelta patches are handled by the caller via the `xdelta3` binary, not here.
"""
import sys, zlib, struct

def die(code, msg):
    sys.stderr.write(msg + "\n"); sys.exit(code)

def read(p):
    with open(p, "rb") as f: return f.read()

# ---------- IPS ----------
def apply_ips(patch, base):
    if patch[:5] != b"PATCH": die(2, "IPS: bad header")
    out = bytearray(base)
    i = 5
    while True:
        rec = patch[i:i+3]
        if rec == b"EOF": break
        if len(rec) < 3: die(2, "IPS: truncated")
        off = (rec[0] << 16) | (rec[1] << 8) | rec[2]; i += 3
        size = (patch[i] << 8) | patch[i+1]; i += 2
        if size == 0:  # RLE run
            rle = (patch[i] << 8) | patch[i+1]; i += 2
            val = patch[i:i+1]; i += 1
            chunk = val * rle
            off2 = off + rle
        else:
            chunk = patch[i:i+size]; i += size
            off2 = off + size
        if off2 > len(out): out.extend(b"\x00" * (off2 - len(out)))
        out[off:off2] = chunk
    # optional truncate extension (3-byte length after EOF)
    tail = patch[i+3:i+6] if patch[i:i+3] == b"EOF" else b""
    if len(tail) == 3:
        newlen = (tail[0] << 16) | (tail[1] << 8) | tail[2]
        out = out[:newlen]
    return bytes(out)

# ---------- UPS ----------
def _ups_num(buf, i):
    val = 0; shift = 1
    while True:
        x = buf[i]; i += 1
        val += (x & 0x7f) * shift
        if x & 0x80: break
        shift <<= 7; val += shift
    return val, i

def apply_ups(patch, base):
    if patch[:4] != b"UPS1": die(2, "UPS: bad header")
    i = 4
    src_size, i = _ups_num(patch, i)
    dst_size, i = _ups_num(patch, i)
    if len(base) != src_size:
        sys.stderr.write(f"UPS: base size {len(base)} != expected {src_size} (continuing)\n")
    out = bytearray(base) + b"\x00" * max(0, dst_size - len(base))
    out = out[:dst_size] if dst_size < len(out) else out
    pos = 0
    body_end = len(patch) - 12  # trailing 3 CRC32s (src,dst,patch)
    while i < body_end:
        rel, i = _ups_num(patch, i)
        pos += rel
        while True:
            b = patch[i]; i += 1
            if pos < len(out): out[pos] ^= b
            pos += 1
            if b == 0: break
    src_crc, dst_crc, _ = struct.unpack("<III", patch[-12:])
    if zlib.crc32(base) & 0xffffffff != src_crc:
        sys.stderr.write("UPS: WARNING source CRC mismatch (wrong base ROM?)\n")
    got = zlib.crc32(bytes(out)) & 0xffffffff
    if got != dst_crc:
        die(2, f"UPS: output CRC {got:08x} != expected {dst_crc:08x}")
    return bytes(out)

# ---------- BPS ----------
def apply_bps(patch, base):
    if patch[:4] != b"BPS1": die(2, "BPS: bad header")
    i = 4
    src_size, i = _ups_num(patch, i)
    dst_size, i = _ups_num(patch, i)
    meta_size, i = _ups_num(patch, i)
    i += meta_size
    src_crc, dst_crc, patch_crc = struct.unpack("<III", patch[-12:])
    if zlib.crc32(base) & 0xffffffff != src_crc:
        die(2, f"BPS: source CRC mismatch — wrong base ROM (have {zlib.crc32(base)&0xffffffff:08x}, need {src_crc:08x})")
    out = bytearray(dst_size)
    out_pos = 0; src_rel = 0; dst_rel = 0
    body_end = len(patch) - 12
    while i < body_end:
        data, i = _ups_num(patch, i)
        action = data & 3; length = (data >> 2) + 1
        if action == 0:  # SourceRead
            out[out_pos:out_pos+length] = base[out_pos:out_pos+length]; out_pos += length
        elif action == 1:  # TargetRead
            out[out_pos:out_pos+length] = patch[i:i+length]; i += length; out_pos += length
        elif action == 2:  # SourceCopy
            off, i = _ups_num(patch, i)
            src_rel += (-(off >> 1) if (off & 1) else (off >> 1))
            for _ in range(length):
                out[out_pos] = base[src_rel]; out_pos += 1; src_rel += 1
        else:  # TargetCopy
            off, i = _ups_num(patch, i)
            dst_rel += (-(off >> 1) if (off & 1) else (off >> 1))
            for _ in range(length):
                out[out_pos] = out[dst_rel]; out_pos += 1; dst_rel += 1
    got = zlib.crc32(bytes(out)) & 0xffffffff
    if got != dst_crc:
        die(2, f"BPS: output CRC {got:08x} != expected {dst_crc:08x}")
    return bytes(out)

def main():
    if len(sys.argv) != 4:
        die(2, "usage: romhack-apply.py <patch> <base_rom> <out_rom>")
    patch, base_path, out_path = sys.argv[1], sys.argv[2], sys.argv[3]
    pb = read(patch); base = read(base_path)
    magic = pb[:5]
    if magic[:4] == b"BPS1": out = apply_bps(pb, base)
    elif magic[:4] == b"UPS1": out = apply_ups(pb, base)
    elif magic == b"PATCH": out = apply_ips(pb, base)
    else: die(3, f"unsupported patch magic {magic!r}")
    with open(out_path, "wb") as f: f.write(out)
    print(f"OK {out_path} ({len(out)} bytes, crc32={zlib.crc32(out)&0xffffffff:08x})")

if __name__ == "__main__":
    main()