import cv2
from PIL import Image
import numpy as np
import time
import os

class CaptureCard:
    def __init__(self, device_index=0, width=1920, height=1080, save_dir="screenshots"):
        """
        初始化采集卡（或摄像头）
        :param device_index: 设备索引号（一般是 0/1/2）
        :param width: 采集宽度
        :param height: 采集高度
        :param save_dir: 截图保存目录
        """
        self.device_index = device_index
        self.width = width
        self.height = height
        self.cap = None
        self.region = None
        self.save_dir = save_dir
        os.makedirs(save_dir, exist_ok=True)

    def open(self):
        """打开采集卡"""
        self.cap = cv2.VideoCapture(self.device_index, cv2.CAP_DSHOW)
        if not self.cap.isOpened():
            self.cap = cv2.VideoCapture(self.device_index)
        if not self.cap.isOpened():
            raise RuntimeError(f"无法打开采集设备 index={self.device_index}")

        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
        print(f"采集卡已打开：{self.width}x{self.height}")

    def close(self):
        """关闭采集卡"""
        if self.cap:
            self.cap.release()
            self.cap = None
            print("采集卡已关闭。")
        cv2.destroyAllWindows()

    def getDesktopImg(self):
        """从采集卡获取一帧图像"""
        if self.cap is None:
            self.open()

        ret, frame = self.cap.read()
        if not ret:
            print("无法从采集卡读取帧")
            return None

        im_opencv = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        if self.region:
            left, top, right, bottom = self.region
            im_opencv = im_opencv[top:bottom, left:right]
        im_PIL = Image.fromarray(im_opencv)
        return [im_opencv, im_PIL]

    def preview(self):
        """实时预览 + 每5秒自动截图"""
        if self.cap is None:
            self.open()

        print("按 'q' 退出实时预览")
        last_capture_time = time.time()
        screenshot_count = 0

        while True:
            ret, frame = self.cap.read()
            if not ret:
                print("无法读取视频帧")
                break

            if self.region:
                left, top, right, bottom = self.region
                frame = frame[top:bottom, left:right]

            # 显示视频
            cv2.imshow("CaptureCard Preview", frame)

            # 每5秒自动截图
            now = time.time()
            if now - last_capture_time >= 5:
                screenshot_count += 1
                filename = os.path.join(self.save_dir, f"screenshot_{screenshot_count}.jpg")
                cv2.imwrite(filename, frame)
                print(f"[截图] 已保存：{filename}")
                last_capture_time = now

            # 按 'q' 退出
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break

        self.close()


if __name__ == "__main__":
    card = CaptureCard(device_index=0, width=1920, height=1080)
    card.preview()