import win32con
import win32api
import win32gui
import win32ui
from mss import mss
import cv2
import numpy as np
import time

ScreenX = win32api.GetSystemMetrics(win32con.SM_CXSCREEN)
ScreenY = win32api.GetSystemMetrics(win32con.SM_CYSCREEN)

ALL_S_CNT = 0
ALL_S_T = 0.0

ScreenShot = mss()


def save_image(img, filename):
    # 保存图像
    cv2.imwrite(filename, img)


def show(img):
        cv2.imshow('window', img)
        cv2.waitKey(1)



def get_all_hwnd():
    hwnd_title_lis = []
    def _callback(hwnd, mouse):
        if win32gui.IsWindow(hwnd) and win32gui.IsWindowEnabled(hwnd) and win32gui.IsWindowVisible(hwnd):
            hwnd_title_lis.append((hwnd, win32gui.GetWindowText(hwnd)))
    win32gui.EnumWindows(_callback, 0)
    return hwnd_title_lis


def screen_shot_mss(width, high, left=None, top=None):
    global ALL_S_T, ALL_S_CNT

    start = time.time()
    if left and top:
        pass
    else:
        left = int(ScreenX / 2 - width/2)
        top = int(ScreenY / 2 - high/2)
    screenshot = ScreenShot.grab((left, top, left+width, top+high))
    img = cv2.cvtColor(np.array(screenshot), cv2.COLOR_BGRA2BGR)

    ALL_S_CNT += 1
    ALL_S_T += (time.time() - start) * 1000
    print(end='\r')
    print(f"平均截图用时: {round(ALL_S_T / ALL_S_CNT, 2)}ms", end="", flush=True)
    return img


def screen_shot_win32(size, title_info):
    global title, ALL_S_T, ALL_S_CNT
    try:
        if title != "":
            hwnd = win32gui.FindWindow(None, title)
            if hwnd == 0:
                print(f'{title}不存在，请重新检查窗口名字是否正确。')
            start = time.time()

            hwndDC = win32gui.GetWindowDC(hwnd)
            mfcDC = win32ui.CreateDCFromHandle(hwndDC)
            saveDC = mfcDC.CreateCompatibleDC()
            saveBitMap = win32ui.CreateBitmap()

            left = int(ScreenX / 2 - size / 2)
            top = int(ScreenY / 2 - size / 2)
            x, y, w, h = (left, top, size, size)
            saveBitMap.CreateCompatibleBitmap(mfcDC, w, h)
            saveDC.SelectObject(saveBitMap)

            saveDC.BitBlt((0, 0), (w, h), mfcDC, (x, y), win32con.SRCCOPY)
            signe_ints_array = saveBitMap.GetBitmapBits(True)
            img = np.frombuffer(signe_ints_array, dtype='uint8')
            img.shape = (h, w, 4)
            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

            win32gui.DeleteObject(saveBitMap.GetHandle())
            saveDC.DeleteDC()
            mfcDC.DeleteDC()

            ALL_S_CNT += 1
            ALL_S_T += (time.time() - start) * 1000
            print(end='\r')
            print(f"平均截图用时: {round(ALL_S_T / ALL_S_CNT, 2)}ms", end="", flush=True)
            return img
    except Exception as e:
        print(e)
        all_title_lis = get_all_hwnd()
        title_lis = []
        for title_item in all_title_lis:
            if title_info in title_item[1]:
                title = title_item[1]
            title_lis.append(title_item[1])
        else:
            print(f"没有找到对应的title, all title: {title_lis}")
            return


if __name__ == '__main__':
    print(f"所有的窗口: {[i[1] for i in get_all_hwnd()]}")
    while True:
        time.sleep(1)
        img_array = screen_shot_mss(25, 1180, 1000)
        # img_array = screen_shot_win32(25, "Client")

        try:
            show(img_array)
            save_image(img_array, f"./tmp/{ALL_S_CNT}.jpg")
        except Exception as e:
            print(e)