#!/usr/bin/env python3
"""查 exe 里到底嵌了哪套图标（PE 资源解析），可导出最大帧为 PNG 核对。

用途：换图标后 Windows 仍显示旧图标时，先确认「exe 里是新图标」还是「外壳图标缓存没刷」。

  python tools/pe_icon.py dist/TfsExcelImporter.exe [out_256.png]
"""

import struct
import sys
import zlib


def u16(d, o):
    return struct.unpack_from("<H", d, o)[0]


def u32(d, o):
    return struct.unpack_from("<I", d, o)[0]


def parse_resources(path):
    d = open(path, "rb").read()
    pe = u32(d, 0x3C)
    if d[pe:pe + 4] != b"PE\0\0":
        raise SystemExit("不是 PE 文件")
    coff = pe + 4
    nsec = u16(d, coff + 2)
    opt_size = u16(d, coff + 16)
    opt = coff + 20
    magic = u16(d, opt)
    dd = opt + (96 if magic == 0x10B else 112)
    rsrc_rva = u32(d, dd + 2 * 8)

    sec = opt + opt_size
    sections = []
    for i in range(nsec):
        s = sec + 40 * i
        vsize, vaddr, rawsize, rawptr = (u32(d, s + 8), u32(d, s + 12),
                                         u32(d, s + 16), u32(d, s + 20))
        sections.append((vaddr, vsize, rawptr, rawsize))

    def rva2off(rva):
        for vaddr, vsize, rawptr, rawsize in sections:
            if vaddr <= rva < vaddr + max(vsize, rawsize):
                return rva - vaddr + rawptr
        raise ValueError("rva %#x 不在任何节里" % rva)

    base = rva2off(rsrc_rva)
    out = []

    def walk(off, path_ids):
        nnamed, nid = u16(d, base + off + 12), u16(d, base + off + 14)
        for i in range(nnamed + nid):
            e = base + off + 16 + 8 * i
            name_id, off_to = u32(d, e), u32(d, e + 4)
            if off_to & 0x80000000:
                walk(off_to & 0x7FFFFFFF, path_ids + [name_id])
            else:
                de = base + off_to
                rva, size = u32(d, de), u32(d, de + 4)
                o = rva2off(rva)
                out.append((path_ids + [name_id], d[o:o + size]))

    walk(0, [])
    return out


def png_write(path, w, h, rgba):
    raw = bytearray()
    for y in range(h):
        raw.append(0)
        raw += rgba[y * w * 4:(y + 1) * w * 4]

    def chunk(tag, data):
        return (struct.pack(">I", len(data)) + tag + data
                + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF))

    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n"
                + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 6, 0, 0, 0))
                + chunk(b"IDAT", zlib.compress(bytes(raw), 9)) + chunk(b"IEND", b""))


def dib_to_rgba(payload):
    bw, bh2 = struct.unpack_from("<ii", payload, 4)
    h = bh2 // 2
    px = payload[40:40 + bw * h * 4]
    out = bytearray(bw * h * 4)
    for y in range(h):
        sy = h - 1 - y
        for x in range(bw):
            si = (sy * bw + x) * 4
            di = (y * bw + x) * 4
            out[di] = px[si + 2]
            out[di + 1] = px[si + 1]
            out[di + 2] = px[si]
            out[di + 3] = px[si + 3]
    return bw, h, out


def main():
    exe = sys.argv[1]
    out_png = sys.argv[2] if len(sys.argv) > 2 else None
    res = parse_resources(exe)

    groups = [(p, b) for p, b in res if p[0] == 14]
    icons = {(p[1] if len(p) > 1 else 0): b for p, b in res if p[0] == 3}
    print("%s: RT_GROUP_ICON=%d RT_ICON=%d" % (exe, len(groups), len(icons)))
    if not groups:
        print("!! exe 里没有图标资源（会用系统默认 exe 图标）")
        return

    for path_ids, b in groups:
        cnt, = struct.unpack_from("<H", b, 4)
        print("  group %s: %d 帧" % (path_ids, cnt))
        frames = []
        for i in range(cnt):
            w, h, cc, resv, planes, bc, sz, idx = struct.unpack_from("<BBBBHHIH", b, 6 + 14 * i)
            frames.append((w or 256, idx, sz))
        print("  尺寸:", [f[0] for f in frames])

        if out_png:
            for size, idx, sz in frames:
                if size == 256:
                    payload = icons.get(idx)
                    if payload and payload[:8] == b"\x89PNG\r\n\x1a\n":
                        open(out_png, "wb").write(payload)
                        print("  已导出 256 PNG 帧 ->", out_png)
                        return
            for size, idx, sz in frames:
                if size == 48:
                    w, h, rgba = dib_to_rgba(icons[idx])
                    png_write(out_png, w, h, rgba)
                    print("  已导出 48 DIB 帧 ->", out_png)
                    return


if __name__ == "__main__":
    main()