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

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"Total game windows: {len(game_windows)}")

# Screen layout: 2x2 grid or tiled
# Monitor resolution is 2560 x 1440 (or 1920x1080)
# Let's check screen width/height
screen_w = user32.GetSystemMetrics(0)
screen_h = user32.GetSystemMetrics(1)
print(f"Screen resolution: {screen_w}x{screen_h}")

coords = [
    (0, 0),
    (805, 0),
    (0, 630),
    (805, 630),
    (1610, 0),
]

for idx, (h, pid) in enumerate(game_windows):
    pos = coords[idx % len(coords)]
    user32.ShowWindow(h, win32con.SW_RESTORE)
    user32.MoveWindow(h, pos[0], pos[1], 806, 629, True)
    print(f"Moved Window {idx} (PID {pid}, HWND 0x{h:X}) to ({pos[0]}, {pos[1]})")

time.sleep(1.0)

for idx, (h, pid) in enumerate(game_windows):
    rect = win32gui.GetWindowRect(h)
    print(f"Window {idx} (PID {pid}) actual rect: {rect}")
    img = ImageGrab.grab(bbox=rect)
    save_path = fr"C:\Projects\JX\logs\tile_game_{idx}_pid_{pid}.png"
    img.save(save_path)
    print(f"Saved tiled screenshot to {save_path}")

# Also capture full desktop overview!
shot = ImageGrab.grab()
shot.save(r"C:\Projects\JX\logs\desktop_all_games.png")
print("Saved desktop overview to logs/desktop_all_games.png")