import ctypes
import sys
import time
from ctypes import wintypes
import win32gui
import win32con
import psutil
from PIL import ImageGrab

sys.stdout.reconfigure(encoding='utf-8')
user32 = ctypes.windll.user32
hdesk = user32.OpenInputDesktop(0, False, 0x01FF)
if hdesk:
    user32.SetThreadDesktop(hdesk)

vggame_pids = [p.info['pid'] for p in psutil.process_iter(['pid', 'name']) if 'vggame' in p.info['name'].lower()]
print(f"Total vggame PIDs: {len(vggame_pids)}")

game_windows = []
def cb(h, _):
    if not user32.IsWindowVisible(h):
        return True
    pid = wintypes.DWORD()
    user32.GetWindowThreadProcessId(h, ctypes.byref(pid))
    if pid.value in vggame_pids:
        cls_buf = ctypes.create_unicode_buffer(256)
        user32.GetClassNameW(h, cls_buf, 256)
        if cls_buf.value == "Sword3 Class":
            game_windows.append((h, pid.value))
    return True

WNDENUMPROC = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HWND, wintypes.LPARAM)
user32.EnumWindows(WNDENUMPROC(cb), 0)

print(f"Found {len(game_windows)} Sword3 Class windows.")

for idx, (h, pid) in enumerate(game_windows):
    user32.ShowWindow(h, win32con.SW_RESTORE)
    user32.SetWindowPos(h, win32con.HWND_TOP, 0, 0, 806, 629, win32con.SWP_SHOWWINDOW)
    user32.SetForegroundWindow(h)
    time.sleep(0.8)
    rect = win32gui.GetWindowRect(h)
    img = ImageGrab.grab(bbox=(0, 0, 806, 629))
    save_path = fr"C:\Projects\JX\logs\vggame_{pid}_{idx}.png"
    img.save(save_path)
    print(f"Saved window {idx} (PID {pid}, HWND 0x{h:X}) to {save_path}")