"""按窗口标题截图。

用法：
    python capture_window.py <title-substring> [output.png]

例：
    python capture_window.py 珠宝壹佰 shot.png
    python capture_window.py 登录 login.png

依赖：Pillow（项目已装），ctypes（标准库）。
仅 Windows。
"""
import ctypes
import sys
from ctypes import wintypes
from pathlib import Path
from PIL import ImageGrab

user32 = ctypes.windll.user32
EnumWindowsProc = ctypes.WINFUNCTYPE(
    ctypes.c_bool, wintypes.HWND, wintypes.LPARAM
)


def _get_window_text(hwnd):
    length = user32.GetWindowTextLengthW(hwnd)
    if length == 0:
        return ""
    buf = ctypes.create_unicode_buffer(length + 1)
    user32.GetWindowTextW(hwnd, buf, length + 1)
    return buf.value


_PROCESS_NAME_CACHE = {}


def _get_process_name(pid: int) -> str:
    """通过 pid 拿到进程可执行文件名（小写，不含路径）。"""
    if pid in _PROCESS_NAME_CACHE:
        return _PROCESS_NAME_CACHE[pid]
    PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
    kernel32 = ctypes.windll.kernel32
    psapi = ctypes.windll.psapi
    h = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
    if not h:
        _PROCESS_NAME_CACHE[pid] = ""
        return ""
    try:
        buf = ctypes.create_unicode_buffer(260)
        size = wintypes.DWORD(260)
        if kernel32.QueryFullProcessImageNameW(h, 0, buf, ctypes.byref(size)):
            name = buf.value.split("\\")[-1].lower()
        else:
            name = ""
    finally:
        kernel32.CloseHandle(h)
    _PROCESS_NAME_CACHE[pid] = name
    return name


def find_window(title_substring: str, exe_filter: tuple = ("pythonw.exe", "python.exe", "zb100imagegenerator.exe")):
    """返回首个标题包含 title_substring 且属于 exe_filter 进程的可见窗口 hwnd。"""
    found = []

    def callback(hwnd, lparam):
        if not user32.IsWindowVisible(hwnd):
            return True
        text = _get_window_text(hwnd)
        if title_substring not in text:
            return True
        # 检查进程
        pid = wintypes.DWORD()
        user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid))
        name = _get_process_name(pid.value)
        if name in exe_filter:
            found.append((hwnd, text))
            return False  # 停止枚举
        return True

    user32.EnumWindows(EnumWindowsProc(callback), 0)
    return found[0] if found else None


def get_window_bbox(hwnd):
    """获取窗口外接矩形（含标题栏 / 边框）。"""
    rect = wintypes.RECT()
    # DWMWA_EXTENDED_FRAME_BOUNDS = 9，能避开 Win10/11 的不可见 shadow border
    DWMWA_EXTENDED_FRAME_BOUNDS = 9
    dwmapi = ctypes.windll.dwmapi
    res = dwmapi.DwmGetWindowAttribute(
        wintypes.HWND(hwnd),
        ctypes.c_uint(DWMWA_EXTENDED_FRAME_BOUNDS),
        ctypes.byref(rect),
        ctypes.sizeof(rect),
    )
    if res != 0:
        # 回退到 GetWindowRect
        user32.GetWindowRect(hwnd, ctypes.byref(rect))
    return (rect.left, rect.top, rect.right, rect.bottom)


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    title = sys.argv[1]
    output = Path(sys.argv[2] if len(sys.argv) > 2 else "shot.png").resolve()

    hit = find_window(title)
    if not hit:
        print(f"找不到标题包含 '{title}' 的窗口", file=sys.stderr)
        sys.exit(2)
    hwnd, full_title = hit

    # 用 PrintWindow 直接从窗口拿位图，不需要把窗口提前（避免打断用户）
    bbox = get_window_bbox(hwnd)
    width = bbox[2] - bbox[0]
    height = bbox[3] - bbox[1]

    img = _print_window_to_pil(hwnd, width, height)
    if img is None:
        # 回退：把窗口提前再截屏
        import time
        user32.ShowWindow(hwnd, 9)  # SW_RESTORE
        user32.SetForegroundWindow(hwnd)
        # 顶置-取消顶置技巧绕开 Windows 前台限制
        HWND_TOPMOST = -1
        HWND_NOTOPMOST = -2
        SWP_NOMOVE = 0x0002
        SWP_NOSIZE = 0x0001
        user32.SetWindowPos(hwnd, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE)
        user32.SetWindowPos(hwnd, HWND_NOTOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE)
        time.sleep(0.5)
        img = ImageGrab.grab(bbox=bbox, all_screens=True)

    img.save(output)
    print(f"OK  hwnd={hwnd}  title={full_title!r}  size={width}x{height}  -> {output}")


def _print_window_to_pil(hwnd, width, height):
    """用 PrintWindow API 从窗口直接拿位图（不依赖窗口在前台）。"""
    from PIL import Image
    gdi32 = ctypes.windll.gdi32

    hdcWindow = user32.GetDC(hwnd)
    hdcMem = gdi32.CreateCompatibleDC(hdcWindow)
    hbm = gdi32.CreateCompatibleBitmap(hdcWindow, width, height)
    gdi32.SelectObject(hdcMem, hbm)

    PW_RENDERFULLCONTENT = 0x00000002
    ok = user32.PrintWindow(hwnd, hdcMem, PW_RENDERFULLCONTENT)

    if ok:
        # 提取位图数据
        class BITMAPINFOHEADER(ctypes.Structure):
            _fields_ = [
                ("biSize", wintypes.DWORD),
                ("biWidth", wintypes.LONG),
                ("biHeight", wintypes.LONG),
                ("biPlanes", wintypes.WORD),
                ("biBitCount", wintypes.WORD),
                ("biCompression", wintypes.DWORD),
                ("biSizeImage", wintypes.DWORD),
                ("biXPelsPerMeter", wintypes.LONG),
                ("biYPelsPerMeter", wintypes.LONG),
                ("biClrUsed", wintypes.DWORD),
                ("biClrImportant", wintypes.DWORD),
            ]

        class BITMAPINFO(ctypes.Structure):
            _fields_ = [("bmiHeader", BITMAPINFOHEADER), ("bmiColors", wintypes.DWORD * 3)]

        bmi = BITMAPINFO()
        bmi.bmiHeader.biSize = ctypes.sizeof(BITMAPINFOHEADER)
        bmi.bmiHeader.biWidth = width
        bmi.bmiHeader.biHeight = -height  # top-down
        bmi.bmiHeader.biPlanes = 1
        bmi.bmiHeader.biBitCount = 32
        bmi.bmiHeader.biCompression = 0  # BI_RGB

        buf_len = width * height * 4
        buf = (ctypes.c_ubyte * buf_len)()
        gdi32.GetDIBits(hdcMem, hbm, 0, height, buf, ctypes.byref(bmi), 0)
        img = Image.frombuffer("RGBA", (width, height), bytes(buf), "raw", "BGRA", 0, 1).convert("RGB")
    else:
        img = None

    gdi32.DeleteObject(hbm)
    gdi32.DeleteDC(hdcMem)
    user32.ReleaseDC(hwnd, hdcWindow)

    return img


if __name__ == "__main__":
    main()