DirtyLeague_Bot/window_utils.py

"""
Utility module for window management, screen capture, and coordinate translation
for the Unity game bot.
"""

import ctypes
from ctypes import wintypes
import time
from typing import Optional, Tuple
import numpy as np
from PIL import Image

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

# Enable Per-Monitor DPI Awareness V2
try:
    shcore = ctypes.windll.shcore
    shcore.SetProcessDpiAwareness(2)
except Exception:
    pass

DESKTOP_ALL_ACCESS = 0x01FF
PW_RENDERFULLCONTENT = 2
SRCCOPY = 0x00CC0020
SW_RESTORE = 9
SW_SHOW = 5


def ensure_interactive_desktop():
    """
    Ensures the current thread is attached to the interactive 'Default' desktop.
    Required when running under background/isolated sandbox sessions.
    """
    try:
        h_desk = user32.OpenDesktopW("Default", 0, False, DESKTOP_ALL_ACCESS)
        if h_desk:
            user32.SetThreadDesktop(h_desk)
    except Exception:
        pass


ensure_interactive_desktop()


class WindowManager:
    """Manages window discovery, activation, and capture."""

    def __init__(self, window_title: str = "DirtyLeague", process_name: str = "DirtyLeague.exe"):
        self.window_title = window_title
        self.process_name = process_name
        self._hwnd: Optional[int] = None

    def find_window(self) -> Optional[int]:
        """Finds the target game window handle."""
        ensure_interactive_desktop()

        # 1. Search by exact window title
        hwnd = user32.FindWindowW(None, self.window_title)
        if hwnd and user32.IsWindow(hwnd):
            self._hwnd = hwnd
            return hwnd

        # 2. Search by partial title or process
        found_hwnds = []
        WNDENUMPROC = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HWND, wintypes.LPARAM)

        def _enum_callback(h, _):
            if user32.IsWindowVisible(h):
                length = user32.GetWindowTextLengthW(h)
                if length > 0:
                    buff = ctypes.create_unicode_buffer(length + 1)
                    user32.GetWindowTextW(h, buff, length + 1)
                    title = buff.value
                    if self.window_title.lower() in title.lower():
                        found_hwnds.append(h)
            return True

        h_desk = user32.OpenDesktopW("Default", 0, False, DESKTOP_ALL_ACCESS)
        if h_desk:
            user32.EnumDesktopWindows(h_desk, WNDENUMPROC(_enum_callback), 0)
            user32.CloseDesktop(h_desk)
        else:
            user32.EnumWindows(WNDENUMPROC(_enum_callback), 0)

        if found_hwnds:
            self._hwnd = found_hwnds[0]
            return self._hwnd

        return None

    @property
    def hwnd(self) -> int:
        """Returns active hwnd or searches for it."""
        if not self._hwnd or not user32.IsWindow(self._hwnd):
            hwnd = self.find_window()
            if not hwnd:
                raise RuntimeError(f"Window '{self.window_title}' not found. Ensure the game is running.")
        return self._hwnd

    def focus(self) -> bool:
        """Brings the window to the foreground and unminimizes if needed."""
        hwnd = self.hwnd
        if user32.IsIconic(hwnd):
            user32.ShowWindow(hwnd, SW_RESTORE)
            time.sleep(0.2)

        user32.ShowWindow(hwnd, SW_SHOW)
        user32.SetForegroundWindow(hwnd)
        time.sleep(0.1)
        return True

    def get_window_rect(self) -> Tuple[int, int, int, int]:
        """Returns (left, top, width, height) of the entire window."""
        rect = wintypes.RECT()
        user32.GetWindowRect(self.hwnd, ctypes.byref(rect))
        return rect.left, rect.top, rect.right - rect.left, rect.bottom - rect.top

    def get_client_rect(self) -> Tuple[int, int, int, int]:
        """Returns (client_left, client_top, client_width, client_height) in screen coordinates."""
        hwnd = self.hwnd
        crect = wintypes.RECT()
        user32.GetClientRect(hwnd, ctypes.byref(crect))
        width = crect.right - crect.left
        height = crect.bottom - crect.top

        pt = wintypes.POINT(0, 0)
        user32.ClientToScreen(hwnd, ctypes.byref(pt))
        return pt.x, pt.y, width, height

    def capture_frame(self, client_only: bool = True) -> np.ndarray:
        """
        Captures the current window content and returns a BGR NumPy array (OpenCV format).
        Uses direct Desktop BitBlt for highest fidelity and hardware acceleration support.
        """
        ensure_interactive_desktop()
        hwnd = self.hwnd

        c_left, c_top, c_width, c_height = self.get_client_rect()
        if c_width <= 0 or c_height <= 0:
            raise RuntimeError(f"Invalid client dimensions: {c_width}x{c_height}")

        if client_only:
            capture_x, capture_y, capture_w, capture_h = c_left, c_top, c_width, c_height
            hdc_source = user32.GetDC(0) # Desktop DC
        else:
            w_left, w_top, w_width, w_height = self.get_window_rect()
            capture_x, capture_y, capture_w, capture_h = w_left, w_top, w_width, w_height
            hdc_source = user32.GetDC(0)

        mem_dc = gdi32.CreateCompatibleDC(hdc_source)
        save_bitmap = gdi32.CreateCompatibleBitmap(hdc_source, capture_w, capture_h)
        gdi32.SelectObject(mem_dc, save_bitmap)

        gdi32.BitBlt(mem_dc, 0, 0, capture_w, capture_h, hdc_source, capture_x, capture_y, SRCCOPY)

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

        bmi = BITMAPINFOHEADER()
        bmi.biSize = ctypes.sizeof(BITMAPINFOHEADER)
        bmi.biWidth = capture_w
        bmi.biHeight = -capture_h
        bmi.biPlanes = 1
        bmi.biBitCount = 32
        bmi.biCompression = 0

        buffer_len = capture_w * capture_h * 4
        buf = (ctypes.c_ubyte * buffer_len)()
        gdi32.GetDIBits(mem_dc, save_bitmap, 0, capture_h, ctypes.byref(buf), ctypes.byref(bmi), 0)

        gdi32.DeleteObject(save_bitmap)
        gdi32.DeleteDC(mem_dc)
        user32.ReleaseDC(0, hdc_source)

        raw_arr = np.ctypeslib.as_array(buf).reshape((capture_h, capture_w, 4))
        return raw_arr[:, :, :3].copy()

    def client_to_screen(self, x: int, y: int) -> Tuple[int, int]:
        """Translates client coordinates (x, y) to absolute desktop screen coordinates."""
        c_left, c_top, _, _ = self.get_client_rect()
        return c_left + x, c_top + y

    def rel_to_client(self, rx: float, ry: float) -> Tuple[int, int]:
        """Translates normalized relative coordinates (0.0..1.0) to client pixel coordinates (x, y)."""
        _, _, width, height = self.get_client_rect()
        return int(rx * width), int(ry * height)

    def rel_to_client_rect(self, rx1: float, ry1: float, rx2: float, ry2: float) -> Tuple[int, int, int, int]:
        """Translates normalized relative ROI (rx1, ry1, rx2, ry2) to client pixel rect (x1, y1, x2, y2)."""
        _, _, width, height = self.get_client_rect()
        return int(rx1 * width), int(ry1 * height), int(rx2 * width), int(ry2 * height)

    def rel_to_screen(self, rx: float, ry: float) -> Tuple[int, int]:
        """Translates normalized relative coordinates (0.0..1.0) directly to absolute screen coordinates."""
        cx, cy = self.rel_to_client(rx, ry)
        return self.client_to_screen(cx, cy)


if __name__ == "__main__":
    wm = WindowManager()
    hwnd = wm.find_window()
    print(f"Found Window HWND: {hwnd}")
    if hwnd:
        rect = wm.get_window_rect()
        crect = wm.get_client_rect()
        print(f"Window Rect: {rect}")
        print(f"Client Rect: {crect}")
        frame = wm.capture_frame(client_only=True)
        print(f"Captured frame shape: {frame.shape}")