#!/usr/bin/env python3
"""冒烟验证：直接拉起 Release exe，截取真实窗口（含标题栏图标），存 PNG。

用途：确认 ApplicationIcon 与 Form.Icon 都生效、且换图标后窗口能正常起来。
纯标准库 + ctypes，不依赖任何第三方包。
"""

import ctypes
import os
import struct
import subprocess
import sys
import time
import zlib
from ctypes import wintypes as wt

u32 = ctypes.WinDLL("user32", use_last_error=True)
g32 = ctypes.WinDLL("gdi32", use_last_error=True)
k32 = ctypes.WinDLL("kernel32", use_last_error=True)

u32.PrintWindow.argtypes = [wt.HWND, wt.HDC, wt.UINT]
u32.GetWindowRect.argtypes = [wt.HWND, ctypes.POINTER(wt.RECT)]
u32.GetWindowDC.argtypes = [wt.HWND]
u32.GetWindowDC.restype = wt.HDC
u32.GetWindowTextW.argtypes = [wt.HWND, wt.LPWSTR, ctypes.c_int]
u32.GetWindowThreadProcessId.argtypes = [wt.HWND, ctypes.POINTER(wt.DWORD)]
u32.IsWindowVisible.argtypes = [wt.HWND]
u32.SendMessageW.argtypes = [wt.HWND, wt.UINT, wt.WPARAM, wt.LPARAM]
u32.SendMessageW.restype = ctypes.c_void_p
g32.CreateCompatibleDC.argtypes = [wt.HDC]
g32.CreateCompatibleDC.restype = wt.HDC
g32.CreateCompatibleBitmap.argtypes = [wt.HDC, ctypes.c_int, ctypes.c_int]
g32.CreateCompatibleBitmap.restype = wt.HBITMAP
g32.SelectObject.argtypes = [wt.HDC, wt.HGDIOBJ]
g32.SelectObject.restype = wt.HGDIOBJ
g32.GetDIBits.argtypes = [wt.HDC, wt.HBITMAP, wt.UINT, wt.UINT, ctypes.c_void_p,
                          ctypes.c_void_p, wt.UINT]
u32.ReleaseDC.argtypes = [wt.HWND, wt.HDC]
u32.ReleaseDC.restype = ctypes.c_int
u32.GetDC.argtypes = [wt.HWND]
u32.GetDC.restype = wt.HDC
u32.FillRect.argtypes = [wt.HDC, ctypes.POINTER(wt.RECT), wt.HBRUSH]
g32.DeleteObject.argtypes = [wt.HGDIOBJ]
g32.DeleteObject.restype = wt.BOOL
g32.DeleteDC.argtypes = [wt.HDC]
g32.DeleteDC.restype = wt.BOOL
g32.CreateSolidBrush.argtypes = [wt.DWORD]
g32.CreateSolidBrush.restype = wt.HBRUSH


class BITMAPINFOHEADER(ctypes.Structure):
    _fields_ = [("biSize", wt.DWORD), ("biWidth", wt.LONG), ("biHeight", wt.LONG),
                ("biPlanes", wt.WORD), ("biBitCount", wt.WORD),
                ("biCompression", wt.DWORD), ("biSizeImage", wt.DWORD),
                ("biXPelsPerMeter", wt.LONG), ("biYPelsPerMeter", wt.LONG),
                ("biClrUsed", wt.DWORD), ("biClrImportant", wt.DWORD)]


def bgra_to_rgba(buf):
    """GetDIBits 给的是 BGRA，PNG 要 RGBA。"""
    out = bytearray(buf)
    for i in range(0, len(out), 4):
        out[i], out[i + 2] = out[i + 2], out[i]
    return bytes(out)


def png(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 find_window(pid, timeout=25.0):
    result = []

    @ctypes.WINFUNCTYPE(wt.BOOL, wt.HWND, wt.LPARAM)
    def cb(hwnd, _):
        p = wt.DWORD()
        u32.GetWindowThreadProcessId(hwnd, ctypes.byref(p))
        if p.value != pid or not u32.IsWindowVisible(hwnd):
            return True
        n = u32.GetWindowTextLengthW(hwnd)
        if n <= 0:
            return True
        buf = ctypes.create_unicode_buffer(n + 2)
        u32.GetWindowTextW(hwnd, buf, n + 2)
        result.append((hwnd, buf.value))
        return False

    end = time.time() + timeout
    while time.time() < end:
        u32.EnumWindows(cb, 0)
        if result:
            return result[0]
        time.sleep(0.3)
    return None


def grab(hwnd, path):
    r = wt.RECT()
    u32.GetWindowRect(hwnd, ctypes.byref(r))
    w, h = r.right - r.left, r.bottom - r.top
    if w <= 0 or h <= 0:
        raise RuntimeError("window rect empty")
    src = u32.GetWindowDC(hwnd)
    mem = g32.CreateCompatibleDC(src)
    bmp = g32.CreateCompatibleBitmap(src, w, h)
    old = g32.SelectObject(mem, bmp)
    ok = u32.PrintWindow(hwnd, mem, 2)  # PW_RENDERFULLCONTENT
    bi = BITMAPINFOHEADER()
    bi.biSize = ctypes.sizeof(BITMAPINFOHEADER)
    bi.biWidth, bi.biHeight = w, -h
    bi.biPlanes, bi.biBitCount, bi.biCompression = 1, 32, 0
    buf = ctypes.create_string_buffer(w * h * 4)
    lines = g32.GetDIBits(mem, bmp, 0, h, buf, ctypes.byref(bi), 0)
    g32.SelectObject(mem, old)
    g32.DeleteObject(bmp)
    g32.DeleteDC(mem)
    u32.ReleaseDC(hwnd, src)
    rgba = bgra_to_rgba(buf.raw)
    png(path, w, h, rgba)
    return w, h, ok, lines, rgba


def crop_zoom(src_png_w, src_png_h, rgba, cw, ch, scale, path, ox=0, oy=0):
    out_w, out_h = cw * scale, ch * scale
    dst = bytearray(out_w * out_h * 4)
    for y in range(out_h):
        sy = oy + y // scale
        for x in range(out_w):
            sx = ox + x // scale
            si = (sy * src_png_w + sx) * 4
            di = (y * out_w + x) * 4
            dst[di:di + 4] = rgba[si:si + 4]
    png(path, out_w, out_h, dst)


def main():
    # 必须声明 DPI 感知：否则 GetWindowRect 返回的是被系统虚拟化（/缩放比）的坐标，
    # 抓出来的图只有窗口左上角一块，高 DPI 下右/下必被裁掉。
    # Win10 起用 shcore 的 per-monitor 感知最可靠，老系统退回 user32.SetProcessDPIAware。
    aware = False
    try:
        ctypes.windll.shcore.SetProcessDpiAwareness(2)   # PROCESS_PER_MONITOR_DPI_AWARE
        aware = True
    except Exception:
        pass
    if not aware:
        try:
            aware = bool(u32.SetProcessDPIAware())
        except Exception:
            pass
    print("dpi aware:", aware)

    exe = sys.argv[1]
    out = sys.argv[2]
    icon_dump = sys.argv[3] if len(sys.argv) > 3 else None

    proc = subprocess.Popen([exe], cwd=os.path.dirname(exe))
    try:
        got = find_window(proc.pid)
        if not got:
            print("FAIL: 没找到窗口（可能是启动崩溃），进程存活 =", proc.poll() is None)
            if proc.poll() is None:
                proc.terminate()
            return 1
        hwnd, title = got
        print("window title:", title)

        # WM_GETICON=0x7F, ICON_BIG=1 / ICON_SMALL=0
        big = u32.SendMessageW(hwnd, 0x7F, 1, 0)
        small = u32.SendMessageW(hwnd, 0x7F, 0, 0)
        print("WM_GETICON big=%#x small=%#x" % (big or 0, small or 0))

        time.sleep(0.6)
        w, h, ok, lines, pixels = grab(hwnd, out)
        print("printed %dx%d ok=%s lines=%d -> %s" % (w, h, ok, lines, out))
        # 标题栏局部放大 4 倍，专门看图标
        crop_zoom(w, h, pixels, min(300, w), min(44, h), 4,
                  os.path.splitext(out)[0] + "_titlebar.png", 0, 0)
        print("titlebar crop ->", os.path.splitext(out)[0] + "_titlebar.png")

        if icon_dump and big:
            u32.DrawIconEx.argtypes = [wt.HDC, ctypes.c_int, ctypes.c_int, wt.HICON,
                                       ctypes.c_int, ctypes.c_int, wt.UINT,
                                       wt.HBRUSH, wt.UINT]
            src = u32.GetDC(None)
            mem = g32.CreateCompatibleDC(src)
            size = 32
            bmp = g32.CreateCompatibleBitmap(src, size, size)
            old = g32.SelectObject(mem, bmp)
            # 白底，避免透明区域看不清
            brush = g32.CreateSolidBrush(0x00FFFFFF)
            rect = wt.RECT(0, 0, size, size)
            u32.FillRect(mem, ctypes.byref(rect), brush)
            u32.DrawIconEx(mem, 0, 0, big, size, size, 0, None, 3)  # DI_NORMAL
            bi = BITMAPINFOHEADER()
            bi.biSize = ctypes.sizeof(BITMAPINFOHEADER)
            bi.biWidth, bi.biHeight = size, -size
            bi.biPlanes, bi.biBitCount, bi.biCompression = 1, 32, 0
            buf2 = ctypes.create_string_buffer(size * size * 4)
            g32.GetDIBits(mem, bmp, 0, size, buf2, ctypes.byref(bi), 0)
            g32.SelectObject(mem, old)
            g32.DeleteObject(bmp)
            g32.DeleteObject(brush)
            g32.DeleteDC(mem)
            u32.ReleaseDC(None, src)
            png(icon_dump, size, size, bgra_to_rgba(buf2.raw))
            print("icon dumped ->", icon_dump)
    finally:
        time.sleep(0.2)
        if proc.poll() is None:
            proc.terminate()
    return 0


if __name__ == "__main__":
    sys.exit(main())