#!/usr/bin/env python3
"""Replace the icon inside an already-built .exe, without rebuilding it.

Useful when a binary you want to keep byte-for-byte was built with a bad icon.
Only the RT_ICON / RT_GROUP_ICON resources are rewritten; no code is touched.

    python tools/patch_exe_icon.py "C:/path/to/App.exe"
    python tools/patch_exe_icon.py "C:/path/to/App.exe" --icon other.ico --no-backup

Sizes up to 64px are embedded as 32-bit BGRA DIBs (what Windows shells have
always understood); 128 and 256 go in PNG-compressed, as Vista+ expects.

Windows only. Requires: pillow.
"""
from __future__ import annotations

import argparse
import ctypes
import shutil
import struct
import sys
from ctypes import wintypes
from io import BytesIO
from pathlib import Path

from PIL import Image

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_ICO = ROOT / "src-tauri" / "icons" / "icon.ico"

RT_ICON = 3
RT_GROUP_ICON = 14
LANG_EN_US = 1033
PNG_FROM = 128  # sizes >= this are stored as PNG, smaller ones as DIB

kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
kernel32.BeginUpdateResourceW.argtypes = [wintypes.LPCWSTR, wintypes.BOOL]
kernel32.BeginUpdateResourceW.restype = wintypes.HANDLE
kernel32.UpdateResourceW.argtypes = [
    wintypes.HANDLE, wintypes.LPCWSTR, wintypes.LPCWSTR,
    wintypes.WORD, wintypes.LPVOID, wintypes.DWORD,
]
kernel32.UpdateResourceW.restype = wintypes.BOOL
kernel32.EndUpdateResourceW.argtypes = [wintypes.HANDLE, wintypes.BOOL]
kernel32.EndUpdateResourceW.restype = wintypes.BOOL


def _int_resource(value):
    """MAKEINTRESOURCE: an integer id is passed where a string pointer is expected."""
    return ctypes.cast(ctypes.c_void_p(value), wintypes.LPCWSTR)


def _check(ok, what):
    if not ok:
        err = ctypes.get_last_error()
        raise OSError(err, f"{what} failed: {ctypes.FormatError(err)}")


def to_dib(img: Image.Image) -> bytes:
    """32bpp bottom-up BGRA DIB plus the 1bpp AND mask, as RT_ICON expects."""
    w, h = img.size
    px = img.convert("RGBA").load()

    xor = bytearray()
    for y in range(h - 1, -1, -1):
        for x in range(w):
            r, g, b, a = px[x, y]
            xor += bytes((b, g, r, a))

    stride = ((w + 31) // 32) * 4
    mask = bytearray()
    for y in range(h - 1, -1, -1):
        row = bytearray(stride)
        for x in range(w):
            if px[x, y][3] == 0:  # 1 = transparent
                row[x // 8] |= 0x80 >> (x % 8)
        mask += row

    header = struct.pack(
        "<IiiHHIIiiII", 40, w, h * 2, 1, 32, 0, len(xor) + len(mask), 0, 0, 0, 0
    )
    return bytes(header + xor + mask)


def to_png(img: Image.Image) -> bytes:
    buf = BytesIO()
    img.save(buf, "PNG", optimize=True)
    return buf.getvalue()


def load_sizes(ico_path: Path):
    """Read every image out of a .ico, smallest first."""
    sizes = sorted(Image.open(ico_path).info["sizes"])
    out = []
    for size in sizes:
        im = Image.open(ico_path)
        im.size = size
        out.append(im.convert("RGBA"))
    return out


def build_group(images, first_id: int, blobs) -> bytes:
    """GRPICONDIR + one GRPICONDIRENTRY per image."""
    data = struct.pack("<HHH", 0, 1, len(images))
    for i, (im, blob) in enumerate(zip(images, blobs)):
        w = 0 if im.width >= 256 else im.width
        h = 0 if im.height >= 256 else im.height
        data += struct.pack("<BBBBHHIH", w, h, 0, 0, 1, 32, len(blob), first_id + i)
    return data


def patch(exe: Path, ico: Path, group_id: int, backup: bool) -> None:
    images = load_sizes(ico)
    blobs = [to_png(im) if im.width >= PNG_FROM else to_dib(im) for im in images]

    if backup:
        bak = exe.with_suffix(exe.suffix + ".bak")
        if bak.exists():
            print(f"  backup already exists, leaving it alone: {bak.name}")
        else:
            shutil.copy2(exe, bak)
            print(f"  backup -> {bak.name}")

    handle = kernel32.BeginUpdateResourceW(str(exe), False)
    if not handle:
        _check(False, "BeginUpdateResource")

    try:
        for i, (im, blob) in enumerate(zip(images, blobs)):
            buf = ctypes.create_string_buffer(blob, len(blob))
            _check(
                kernel32.UpdateResourceW(
                    handle, _int_resource(RT_ICON), _int_resource(i + 1),
                    LANG_EN_US, ctypes.cast(buf, wintypes.LPVOID), len(blob),
                ),
                f"UpdateResource RT_ICON #{i + 1}",
            )
            kind = "png" if im.width >= PNG_FROM else "dib"
            print(f"  RT_ICON #{i + 1:<2} {im.width:>3}x{im.height:<3} {kind} {len(blob):>7,}b")

        group = build_group(images, 1, blobs)
        buf = ctypes.create_string_buffer(group, len(group))
        _check(
            kernel32.UpdateResourceW(
                handle, _int_resource(RT_GROUP_ICON), _int_resource(group_id),
                LANG_EN_US, ctypes.cast(buf, wintypes.LPVOID), len(group),
            ),
            "UpdateResource RT_GROUP_ICON",
        )
        print(f"  RT_GROUP_ICON #{group_id} {len(images)} entries")
    except Exception:
        kernel32.EndUpdateResourceW(handle, True)  # discard
        raise

    _check(kernel32.EndUpdateResourceW(handle, False), "EndUpdateResource")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("exe", type=Path)
    ap.add_argument("--icon", type=Path, default=DEFAULT_ICO)
    ap.add_argument("--group-id", type=int, default=32512,
                    help="RT_GROUP_ICON id the shell shows (default: 32512)")
    ap.add_argument("--no-backup", action="store_true")
    args = ap.parse_args()

    if sys.platform != "win32":
        print("This tool uses the Windows resource API.", file=sys.stderr)
        return 1
    for p in (args.exe, args.icon):
        if not p.exists():
            print(f"not found: {p}", file=sys.stderr)
            return 1

    print(f"{args.exe}")
    patch(args.exe, args.icon, args.group_id, not args.no_backup)
    print("done - Explorer caches icons, so run `ie4uinit.exe -show` if the old one lingers")
    return 0


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