import cv2


class Camera:
    """摄像头初始化与帧读取（支持颜色参数）"""

    def __init__(self, cam_id=0, target_color="red"):

        self.cap = cv2.VideoCapture(cam_id)
        self.is_opened = self.cap.isOpened()
        if not self.is_opened:
            raise RuntimeError("无法打开摄像头！")

        self.target_color = target_color.lower()
        if self.target_color not in ["red", "blue"]:
            raise ValueError("仅支持 'red' 或 'blue'")

        self._set_cam_params()

    def _set_cam_params(self):
        """根据颜色设置摄像头参数"""
        self.cap.set(cv2.CAP_PROP_AUTOFOCUS, 0)  # 关闭自动对焦
        self.cap.set(cv2.CAP_PROP_AUTO_WB, 0)    # 关闭自动白平衡
        self.cap.set(cv2.CAP_PROP_ISO_SPEED, 0)  # 自动ISO

        if self.target_color == "red":
            self.cap.set(cv2.CAP_PROP_CONTRAST, 10)
            self.cap.set(cv2.CAP_PROP_EXPOSURE, -10)
            self.cap.set(cv2.CAP_PROP_WB_TEMPERATURE, 5)
        elif self.target_color == "blue":
            self.cap.set(cv2.CAP_PROP_CONTRAST, 10)
            self.cap.set(cv2.CAP_PROP_EXPOSURE, -9)
            self.cap.set(cv2.CAP_PROP_WB_TEMPERATURE, 5)

    def read_frame(self):
        """读取一帧"""
        if not self.is_opened:
            return False, None
        return self.cap.read()

    def release(self):
        """释放资源"""
        if self.is_opened:
            self.cap.release()
            self.is_opened = False

    def switch_color(self, target_color):
        """切换目标颜色并更新参数"""
        target_color = target_color.lower()
        if target_color in ["red", "blue"] and target_color != self.target_color:
            self.target_color = target_color
            self._set_cam_params()
            return True
        return False