#!/usr/bin/env python3
"""apply_patch.py -- DiKon firmware patcher for the Kodak Diconix 150 (RND2 Labs / TheRND2)

Applies an RND2 Labs DiKon IPS patch to a stock Kodak Diconix 150 ROM dump you supply,
and writes both a 32 KB image (for a 27C256) and the 64 KB doubled image
(for a W27C512 / 27C512 drop-in).

Targets:
  v5i  (default) -- the shipping DiKon firmware: v5u code (self-test DW fix) plus the
                    DiKon self-test banner clipped below 0x3000. RECOMMENDED burn.
                    (v5i replaces v5f AND v5i, both WITHDRAWN: v5f overwrote
                    0x3000-0x3018 -- live resident-font glyph cells on real
                    hardware -- and bricked the printer. See CHANGELOG.)
  v5             -- the plain v5 PCL raster firmware, stock Kodak self-test page.
  v5t            -- v5 plus the older RND2 Labs self-test page (superseded by v5i;
                    kept for reproducibility -- see CHANGELOG).

Input ROM may be either:
  - the stock 32K dump  (MD5 a258fa60ac1611b885fd32aad2dd2b58), or
  - a 64K dump of a doubled stock ROM (two identical 32K halves).

The expected output MD5s are built in; the script REFUSES to write anything if
any checksum is wrong. Python 3 standard library only -- no dependencies.

Usage:
  python apply_patch.py STOCK_ROM.bin                 # v5i (default, recommended)
  python apply_patch.py STOCK_ROM.bin --target v5     # plain v5, stock self-test page
  python apply_patch.py STOCK_ROM.bin --target v5i -o out/
  python apply_patch.py STOCK_ROM.bin --ips some.ips  # explicit patch file

Project: thernd2.com/diconix
"""

import argparse
import hashlib
import os
import sys

MD5_STOCK_32K = "a258fa60ac1611b885fd32aad2dd2b58"

# target -> metadata (default ips filename, out basename, expected MD5s)
TARGETS = {
    "v5i": {
        "ips": "d150_dikon_v5i.ips",
        "base": "d150_dikon_v5i",
        "md5_32k": "03a386608973e1c686d49d1ef0439410",
        "md5_64k": "2d0d67cc221731332a82db06d4a3da2d",
        "desc": "DiKon v5i firmware (v5u code + DiKon banner, metal-proven ceiling + font-variety table: 6 distinct self-test rows) -- RECOMMENDED",
    },
    "v5h": {
        "ips": "d150_dikon_v5h.ips",
        "base": "d150_dikon_v5h",
        "md5_32k": "2acd4a5674f6780a4800a8f9c743491a",
        "md5_64k": "1e6dd3469421a69eb1b74c37e9d65335",
        "desc": "DiKon v5h (font-force: all diag rows one font; metal-proven safe)",
    },
    "v5": {
        "ips": "d150_m7v5.ips",
        "base": "d150_m7v5",
        "md5_32k": "fe9215abec1f8c4f4f407867302fa898",
        "md5_64k": "2f2a10f449f522df2500effe616f5f67",
        "desc": "M7 v5 firmware (stock self-test page)",
    },
    "v5t": {
        "ips": "d150_m7v5t.ips",
        "base": "d150_m7v5t",
        "md5_32k": "0f5b5b97c71da971344171f2099df137",
        "md5_64k": "e90ae2cfd7180c96ed15003a5f794868",
        "desc": "M7 v5t firmware (v5 + old RND2 Labs self-test page; superseded by v5i)",
    },
}


def md5(data: bytes) -> str:
    return hashlib.md5(data).hexdigest()


def die(msg: str) -> None:
    sys.exit("ERROR: " + msg)


def load_stock(path: str) -> bytes:
    with open(path, "rb") as f:
        data = f.read()
    if len(data) == 65536:
        if data[:32768] != data[32768:]:
            die("64K input is not a doubled 32K image (halves differ). "
                "Is this really a stock Diconix 150 dump?")
        print("Input is a 64K doubled image; using first 32K half.")
        data = data[:32768]
    elif len(data) != 32768:
        die(f"input is {len(data)} bytes; expected 32768 (27C256) "
            "or 65536 (doubled, W27C512).")
    h = md5(data)
    if h != MD5_STOCK_32K:
        die(f"stock ROM MD5 mismatch.\n  got      {h}\n  expected {MD5_STOCK_32K}\n"
            "This is not the stock Kodak Diconix 150 ROM these patches were built "
            "against. (Diconix 150 Plus / 180si / 300w are different code bases -- "
            "see docs/08-portability.md.)")
    print(f"Stock ROM OK  (MD5 {h})")
    return data


def apply_ips(base: bytes, ips: bytes) -> bytearray:
    if ips[:5] != b"PATCH":
        die("patch file is not an IPS patch (missing PATCH header).")
    out = bytearray(base)
    i = 5
    n = len(ips)
    while True:
        if i + 3 > n:
            die("truncated IPS patch (no EOF marker).")
        tag = ips[i:i + 3]
        if tag == b"EOF" and i + 3 == n:
            break
        off = int.from_bytes(tag, "big")
        i += 3
        size = int.from_bytes(ips[i:i + 2], "big")
        i += 2
        if size == 0:  # RLE record
            run = int.from_bytes(ips[i:i + 2], "big")
            i += 2
            out[off:off + run] = bytes([ips[i]]) * run
            i += 1
        else:
            out[off:off + size] = ips[i:i + size]
            i += size
    return out


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Apply an RND2 Labs DiKon IPS patch to a stock Diconix 150 ROM.")
    ap.add_argument("stock", help="stock ROM dump (32K, or 64K doubled)")
    ap.add_argument("--target", choices=sorted(TARGETS), default="v5i",
                    help="which firmware to build (default: v5i)")
    ap.add_argument("--ips", default=None,
                    help="explicit IPS patch file (default: the target's .ips "
                         "next to this script)")
    ap.add_argument("-o", "--outdir", default=".",
                    help="output directory (default: .)")
    args = ap.parse_args()

    t = TARGETS[args.target]
    here = os.path.dirname(os.path.abspath(__file__))
    ips_path = args.ips or os.path.join(here, t["ips"])
    if not os.path.isfile(ips_path):
        die(f"patch file not found: {ips_path}")

    print(f"Target: {args.target} -- {t['desc']}")
    base = load_stock(args.stock)
    with open(ips_path, "rb") as f:
        ips = f.read()

    patched = bytes(apply_ips(base, ips))
    if len(patched) != 32768:
        die("patched image is not 32768 bytes -- wrong patch file?")

    h32 = md5(patched)
    if h32 != t["md5_32k"]:
        die(f"patched 32K MD5 mismatch.\n  got      {h32}\n  expected "
            f"{t['md5_32k']}\nWrong patch file for target {args.target}?")
    doubled = patched * 2
    h64 = md5(doubled)
    if h64 != t["md5_64k"]:
        die(f"doubled 64K MD5 mismatch.\n  got      {h64}\n  expected "
            f"{t['md5_64k']}")

    os.makedirs(args.outdir, exist_ok=True)
    p32 = os.path.join(args.outdir, t["base"] + ".bin")
    p64 = os.path.join(args.outdir, t["base"] + "_W27C512.bin")
    with open(p32, "wb") as f:
        f.write(patched)
    with open(p64, "wb") as f:
        f.write(doubled)

    print(f"Wrote {p32}  (32768 bytes, MD5 {h32})  -- 27C256")
    print(f"Wrote {p64}  (65536 bytes, MD5 {h64})  -- W27C512 (image doubled)")
    print("All checksums verified. Done.")


if __name__ == "__main__":
    main()
