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)

# All 5 game windows
game_pids = [3268, 19076, 19316, 21200, 27800]
game_hwnds = {}

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

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

print(f"Found {len(game_hwnds)} game HWNDs:")
for pid, h in game_hwnds.items():
    print(f"PID {pid} -> HWND 0x{h:X}")

# Capture each window by bringing it to front at (100, 100)
for pid, h in game_hwnds.items():
    user32.ShowWindow(h, win32con.SW_RESTORE)
    user32.SetWindowPos(h, win32con.HWND_TOPMOST, 50, 50, 806, 629, win32con.SWP_SHOWWINDOW)
    user32.SetForegroundWindow(h)
    time.sleep(0.6)
    
    # Grab from (50, 50, 856, 679)
    img = ImageGrab.grab(bbox=(50, 50, 856, 679))
    save_path = fr"C:\Projects\JX\logs\clean_game_pid_{pid}.png"
    img.save(save_path)
    print(f"Saved clean capture for PID {pid} to {save_path}")
    
    # Restore non-topmost
    user32.SetWindowPos(h, win32con.HWND_NOTOPMOST, 50, 50, 806, 629, win32con.SWP_SHOWWINDOW)