import cv2
import numpy as np
from PIL import Image
from win32api import GetSystemMetrics
import ctypes
import os
import time


# 定义BITMAPINFOHEADER的结构
class BITMAPINFOHEADER(ctypes.Structure):
    _fields_ = [
        ('biSize', ctypes.c_ulong),
        ('biWidth', ctypes.c_long),
        ('biHeight', ctypes.c_long),
        ('biPlanes', ctypes.c_ushort),
        ('biBitCount', ctypes.c_ushort),
        ('biCompression', ctypes.c_ulong),
        ('biSizeImage', ctypes.c_ulong),
        ('biXPelsPerMeter', ctypes.c_long),
        ('biYPelsPerMeter', ctypes.c_long),
        ('biClrUsed', ctypes.c_ulong),
        ('biClrImportant', ctypes.c_ulong)
    ]


class ScreenCapture:

    def __init__(self, output_dir='screenshots', use_key_frame_detection=True, threshold=100000):
        """初始化函数，设定是否使用关键帧检测以及阈值和保存图片的目录"""
        self.output_dir = output_dir
        self.use_key_frame_detection = use_key_frame_detection
        self.threshold = threshold
        self.last_image = None

        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)


    def capture(self):
            """使用win32api和ctypes捕捉屏幕内容"""
            width = GetSystemMetrics(0)
            height = GetSystemMetrics(1)
            hwin = ctypes.windll.user32.GetDesktopWindow()

            hwindc = ctypes.windll.user32.GetWindowDC(hwin)
            srcdc = ctypes.windll.gdi32.CreateCompatibleDC(hwindc)
            srcbmp = ctypes.windll.gdi32.CreateCompatibleBitmap(hwindc, width, height)
            ctypes.windll.gdi32.SelectObject(srcdc, srcbmp)

            ctypes.windll.gdi32.BitBlt(srcdc, 0, 0, width, height, hwindc, 0, 0, 0x00CC0020)

            bmpinfo = BITMAPINFOHEADER()
            bmpinfo.biSize = ctypes.sizeof(BITMAPINFOHEADER)
            bmpinfo.biWidth = width
            bmpinfo.biHeight = height
            bmpinfo.biPlanes = 1  # must be 1
            bmpinfo.biBitCount = 32  # 4 bytes per pixel
            bmpinfo.biCompression = 0  # BI_RGB, no compression
            bmpinfo.biSizeImage = width * height * 4  # size of the image in bytes

            bmpstr = ctypes.create_string_buffer(bmpinfo.biSizeImage)
            ctypes.windll.gdi32.GetDIBits(srcdc, srcbmp, 0, height, bmpstr, ctypes.byref(bmpinfo), 0)

            ctypes.windll.gdi32.DeleteObject(srcbmp)
            ctypes.windll.user32.ReleaseDC(hwin, hwindc)
            ctypes.windll.gdi32.DeleteDC(srcdc)

            image = Image.frombuffer(
                'RGB',
                (bmpinfo.biWidth, bmpinfo.biHeight),
                bmpstr,
                'raw',
                'BGRX',
                0,
                -1
            )
            image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
            return np.array(image)

    def capture_and_check_changes(self):
        """捕获屏幕并检查是否与上一个图像有显著的差异"""
        current_image = self.capture()

        if self.last_image is None:
            self.last_image = current_image
            return current_image, True

        diff = cv2.absdiff(self.last_image, current_image)
        gray_diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)

        has_changes = cv2.countNonZero(gray_diff) > self.threshold
        return current_image, has_changes

    def save_image(self, image):
        """保存图片"""
        cv2.imwrite(f'{self.output_dir}/screenshot_{time.time()}.png', image)

    def start_capture(self):
        """执行一次捕获并判断是否需要保存的过程"""
        current_image, has_changes = self.capture_and_check_changes()

        if not self.use_key_frame_detection or has_changes:
            self.save_image(current_image)

        self.last_image = current_image


# 使用示例：
# 创建ScreenCapture对象，并指定是否使用关键帧检测，以及阈值
capturer = ScreenCapture(output_dir='screenshots', use_key_frame_detection=True, threshold=100000)
# 执行一次捕获并判断是否需要保存的过程
while(True):
    capturer.start_capture()