import time
from ctypes import wintypes
import ctypes
import win32gui
import win32con
import win32ui
import cv2
import numpy as np


class WindowManager:
    @staticmethod
    def Capture_Win32_Gpu(handle: int, model: int = 3):
        """
        通过win32方式截图，并且无视硬件加速
        https://stackoverflow.com/questions/19695214/screenshot-of-inactive-window-printwindow-win32gui
        """
        ctypes.windll.user32.SetProcessDPIAware()  # 抑制缩放

        rect = win32gui.GetWindowRect(handle)
        width, height = rect[2] - rect[0], rect[3] - rect[1]

        hwnd_dc = win32gui.GetWindowDC(handle)
        mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc)
        save_dc = mfc_dc.CreateCompatibleDC()
        save_bit_map = win32ui.CreateBitmap()

        save_bit_map.CreateCompatibleBitmap(mfc_dc, width, height)
        save_dc.SelectObject(save_bit_map)

        ctypes.windll.user32.PrintWindow(handle, save_dc.GetSafeHdc(), model)  # 如果仍然是黑屏，尝试修改数字 ‘3’ (从0开始...)
        bmpinfo = save_bit_map.GetInfo()
        bmpstr = save_bit_map.GetBitmapBits(True)

        capture = np.frombuffer(bmpstr, dtype=np.uint8).reshape((bmpinfo["bmHeight"], bmpinfo["bmWidth"], 4))
        capture = np.ascontiguousarray(capture)[..., :-1]

        win32gui.DeleteObject(save_bit_map.GetHandle())
        save_dc.DeleteDC()
        mfc_dc.DeleteDC()
        win32gui.ReleaseDC(handle, hwnd_dc)

        capture = cv2.cvtColor(capture, cv2.COLOR_RGBA2RGB)
        return capture

    @staticmethod
    def Get_Handle():
        """
        取窗口句柄
        :return: int,窗口句柄
        """
        return win32gui.GetForegroundWindow()

    @staticmethod
    def Get_Window_Position_And_Size(handle):
        """
        窗口取位置和大小
        :param handle: int,窗口句柄
        :return:
        """
        return win32gui.GetWindowRect(handle)

    class WindowRect(ctypes.Structure):
        _fields_ = [
            ("left", ctypes.c_long),
            ("top", ctypes.c_long),
            ("right", ctypes.c_long),
            ("bottom", ctypes.c_long)
        ]

    def __init__(self):
        """
        调用Windows官方API实现对window的操作
        """
        self.user32 = ctypes.windll.user32
        self.kernel32 = ctypes.windll.kernel32
        # self.GetLastError = ctypes.windll.kernel32.GetLastError()
        self.gdi32 = ctypes.WinDLL("Gdi32.dll")

    def Find_Windows(self, parent_handle: int = None, class_name: str = None, window_title: str = None):
        """
        模糊遍历窗口
        :param parent_handle: str,父窗口句柄
        :param class_name: str,欲寻找窗口类名
        :param window_title: str,欲寻找窗口标题
        :return: list,窗口句柄列表
        """
        handle_list = []
        if parent_handle is None:
            parent_handle = self.user32.GetDesktopWindow()
        handle = self.user32.GetWindow(parent_handle, 5)
        while handle != 0:
            valid = True
            if class_name is not None:
                actual_class_name = str(self.Get_Class_Name(handle)).lower()
                if actual_class_name.find(class_name.lower()) == -1:
                    valid = False
                time.sleep(0)
            if window_title is not None:
                actual_window_title = str(self.Get_Window_Title(handle)).lower()
                if actual_window_title.find(window_title.lower()) == -1:
                    valid = False
                time.sleep(0)
            if valid:
                handle_list.append(handle)
            handle = self.user32.GetWindow(handle, 2)
        return handle_list

    def Get_Class_Name(self, handle: int):
        """
        取窗口类名
        :param handle: int,窗口句柄
        :return: str,窗口类名
        """
        GetClassNameA = self.user32.GetClassNameA
        GetClassNameA.argtypes = [ctypes.c_int, ctypes.c_char_p, ctypes.c_int]
        GetClassNameA.restype = ctypes.c_int

        buffer_size = 256
        name_buffer = ctypes.create_string_buffer(buffer_size)
        result = GetClassNameA(handle, name_buffer, buffer_size)
        if result != 0:
            class_name = name_buffer.value.decode('utf-8')
            return class_name
        else:
            return ""

    def Get_Window_Title(self, handle: int):
        """
        取窗口标题
        :param handle: int,窗口句柄
        :return: str,窗口标题
        """
        GetWindowTextLengthW = self.user32.GetWindowTextLengthW
        GetWindowTextLengthW.argtypes = [ctypes.c_int]
        GetWindowTextLengthW.restype = ctypes.c_int

        GetWindowTextW = self.user32.GetWindowTextW
        GetWindowTextW.argtypes = [ctypes.c_int, ctypes.c_wchar_p, ctypes.c_int]
        GetWindowTextW.restype = ctypes.c_int

        length = GetWindowTextLengthW(handle)
        if length > 0:
            buffer_size = length + 1
            title_buffer = ctypes.create_unicode_buffer(buffer_size)
            result = GetWindowTextW(handle, title_buffer, buffer_size)
            if result != 0:
                title = title_buffer.value
                return title
            else:
                return ""
        else:
            return ""

    def Is_Window_Visible(self, handle: int):
        """
        窗口是否可见
        :param handle: int,窗口句柄
        :return: bool,如果该函数成功，则返回值为非零值。如果函数失败，则返回值为零。
        """
        return self.user32.IsWindowVisible(handle)

    def Set_Window_Position_And_Size(self, handle, left=None, top=None, width=None, height=None):
        """
        设置窗口位置和大小
        :param handle: int,窗口句柄
        :param left: int,窗口左侧的新位置
        :param top: int,窗口顶部的新位置
        :param width: int,窗口的新宽度
        :param height: int,窗口的新高度
        :return: bool,如果该函数成功，则返回值为非零值。如果函数失败，则返回值为零
        """
        rect = self.WindowRect()
        self.user32.GetWindowRect(handle, ctypes.byref(rect))
        if left is None:
            left = rect.left
        if top is None:
            top = rect.top
        if width is None:
            width = rect.right - rect.left
        if height is None:
            height = rect.bottom - rect.top
        return self.user32.MoveWindow(handle, left, top, width, height, True)

    def Set_Window_Topmost(self, handle, activate_window=True):
        """
        窗口置顶
        :param handle: int,窗口句柄
        :param activate_window: bool,是否激活窗口
        :return: bool,如果该函数成功，则返回值为非零值。如果函数失败，则返回值为零。
        """
        HWND_TOPMOST = -1
        SWP_SHOWWINDOW = 0x40
        SWP_NOSIZE = 0x1
        SWP_NOMOVE = 0x2
        SWP_NOACTIVATE = 0x10

        flags = SWP_SHOWWINDOW | SWP_NOSIZE | SWP_NOMOVE
        if not activate_window:
            flags |= SWP_NOACTIVATE
        result = self.user32.SetWindowPos(handle, HWND_TOPMOST, 0, 0, 0, 0, flags)
        return result != 0

    def Set_Window_Display_Affinity(self, handle, affinity):
        """
        窗口反截图，仅支持自身调用
        :param handle: int,顶级窗口的句柄。 窗口必须属于当前进程。
        :param affinity: int，0:取消设置;1:黑框;2:透明
        :return: bool,如果该函数成功，则返回值为非零值。如果函数失败，则返回值为零。
        """
        HWND = ctypes.c_void_p
        DWORD = ctypes.c_ulong
        BOOL = ctypes.c_int

        self.user32.SetWindowDisplayAffinity.argtypes = (HWND, DWORD)
        self.user32.SetWindowDisplayAffinity.restype = BOOL

        AFFINITY_CONSTANTS = {
            0: 0x00000000,  # WDA_NONE
            1: 0x00000001,  # WDA_MONITOR
            2: 0x00000011  # WDA_EXCLUDEFROMCAPTURE
        }

        affinity_value = AFFINITY_CONSTANTS.get(affinity, 0)

        result = self.user32.SetWindowDisplayAffinity(handle, affinity_value)
        print(self.kernel32.GetLastError())

        return result

    def Enum_Child_Window_Handles(self, hwnd):
        """
        枚举给定主窗口句柄的子窗口句柄，并返回一个包含子窗口句柄的列表。

        :param hwnd: int，主窗口句柄
        :return: list，子窗口句柄列表
        """
        # 定义Windows API函数的参数类型
        EnumChildProc = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HWND, wintypes.LPARAM)
        EnumChildWindows = self.user32.EnumChildWindows

        # 定义回调函数，用于处理枚举到的子窗口
        def enum_child_proc(hwnd, lparam):
            # 将子窗口句柄添加到列表中
            ctypes.cast(lparam, ctypes.POINTER(ctypes.py_object)).contents.value.append(hwnd)
            return True

        # 创建一个空的子窗口句柄列表
        child_handles = []

        # 枚举子窗口句柄
        EnumChildWindows(hwnd, EnumChildProc(enum_child_proc), ctypes.byref(ctypes.py_object(child_handles)))

        # 返回子窗口句柄列表
        return child_handles

    def Window_Draw_Border_Rectangle(self, handle: int, BorderThickness: int = 2, BorderColor: int = 0xFF0000,
                                     rect_left: int = 0, rect_top: int = 0, rect_right: int = 100,
                                     rect_bottom: int = 100, isTransparent: bool = False):
        """
        窗口绘制矩形
        :param handle: int,窗口句柄
        :param BorderThickness: int,矩形边框粗细
        :param BorderColor: int,矩形边框颜色
        :param rect_left: int,左边
        :param rect_top: int,顶边
        :param rect_right: int,右边
        :param rect_bottom: int,底边
        :param isTransparent: bool,是否透明
        :return:
        """
        # 获取窗口DC
        hdc = win32gui.GetDC(handle)
        # 创建画笔
        pen = win32gui.CreatePen(win32con.PS_SOLID, BorderThickness, BorderColor)
        # 创建画刷
        if isTransparent:
            brush = win32gui.GetStockObject(win32con.HOLLOW_BRUSH)  # 空画刷，不填充
            # 将背景模式设置为透明
            win32gui.SetBkMode(hdc, win32con.TRANSPARENT)
        else:
            brush = win32gui.CreateSolidBrush(0x00000000)  # 0x00000000 表示透明
        try:
            # 使用上下文管理器管理画笔和画刷
            with self._SelectObjects(hdc, pen, brush):
                # 绘制矩形
                result = self.gdi32.Rectangle(hdc, rect_left, rect_top, rect_right, rect_bottom)
        finally:
            # 释放窗口DC
            win32gui.ReleaseDC(handle, hdc)
        return result

    # 自定义上下文管理器来管理GDI对象
    class _SelectObjects:
        def __init__(self, hdc, *objects):
            self.hdc = hdc
            self.objects = objects

        def __enter__(self):
            for obj in self.objects:
                win32gui.SelectObject(self.hdc, obj)
            return self

        def __exit__(self, exc_type, exc_value, traceback):
            for obj in self.objects:
                win32gui.DeleteObject(obj)


if __name__ == "__main__":
    # 示例用法：
    window_manager = WindowManager()

    # 查找所有可见窗口
    # visible_windows = window_manager.Find_Windows()
    # for handle in visible_windows:
    #     class_name = window_manager.Get_Class_Name(handle)
    #     window_title = window_manager.Get_Window_Title(handle)
    #     print(f"Handle: {handle}, Class Name: {class_name}, Window Title: {window_title}")

    # 设置窗口置顶
    # window_handle = 6556218  # 替换为实际窗口句柄
    # window_manager.Set_Window_Topmost(window_handle)

    # 设置窗口位置和大小
    # window_manager.Set_Window_Position_And_Size(window_handle, left=100, top=100, width=800, height=600)
    # 示例用法
    # handle = 1509402  # 获取当前窗口句柄
    # RectangleThickness = 2  # 矩形边框线宽
    # RectangleColor = 0xFF0000  # 矩形颜色 (红色)

    # 指定矩形的坐标和大小
    # rect_left = 100
    # rect_top = 100
    # rect_right = 200
    # rect_bottom = 200

    # 绘制矩形
    # result = window_manager.Window_Draw_Border_Rectangle(handle, RectangleThickness, RectangleColor, rect_left,
    #                                                      rect_top,
    #                                                      rect_right,
    #                                                      rect_bottom, True)