"""Capture timed screenshots of a btl4 window via PrintWindow (works occluded).

Usage: python fxshot.py [count] [interval_s] [outprefix] [nodesubstr]
"""
import ctypes, ctypes.wintypes as wt, sys, time

user32 = ctypes.windll.user32
gdi32 = ctypes.windll.gdi32

count = int(sys.argv[1]) if len(sys.argv) > 1 else 8
interval = float(sys.argv[2]) if len(sys.argv) > 2 else 0.4
prefix = sys.argv[3] if len(sys.argv) > 3 else 'fx'
want = sys.argv[4] if len(sys.argv) > 4 else '1501'

titles = []
def enum_cb(hwnd, lparam):
    n = user32.GetWindowTextLengthW(hwnd)
    if n > 0:
        buf = ctypes.create_unicode_buffer(n + 1)
        user32.GetWindowTextW(hwnd, buf, n + 1)
        if buf.value.startswith('BattleTech') and want in buf.value:
            titles.append((hwnd, buf.value))
    return True
WNDENUMPROC = ctypes.WINFUNCTYPE(ctypes.c_bool, wt.HWND, wt.LPARAM)
user32.EnumWindows(WNDENUMPROC(enum_cb), 0)
if not titles:
    print('no matching BattleTech window'); sys.exit(1)
hwnd, title = titles[0]
print('capturing:', title)

rect = wt.RECT()
user32.GetClientRect(hwnd, ctypes.byref(rect))
w, h = rect.right, rect.bottom

hdc = user32.GetDC(0)
mem = gdi32.CreateCompatibleDC(hdc)
bmp = gdi32.CreateCompatibleBitmap(hdc, w, h)
gdi32.SelectObject(mem, bmp)

class BMIH(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)]

from PIL import Image
for i in range(count):
    # PW_RENDERFULLCONTENT (2) captures D3D content
    user32.PrintWindow(hwnd, mem, 2 | 1)   # +PW_CLIENTONLY
    bmi = BMIH(ctypes.sizeof(BMIH), w, -h, 1, 32, 0, 0, 0, 0, 0, 0)
    buf = ctypes.create_string_buffer(w * h * 4)
    gdi32.GetDIBits(mem, bmp, 0, h, buf, ctypes.byref(bmi), 0)
    img = Image.frombuffer('RGB', (w, h), buf.raw, 'raw', 'BGRX', 0, 1)
    img.save(r'C:\git\bt411\scratchpad\%s_%d.png' % (prefix, i))
    time.sleep(interval)
print('saved %d shots %dx%d' % (count, w, h))
