import depthai as dai
import numpy as np
from camera.base_camera import Camera

class OakCamera(Camera):
    def __init__(self, resolution=(640, 480), fps=30, mirror=False):
        self.pipeline = dai.Pipeline()
        cam = self.pipeline.create(dai.node.ColorCamera)
        cam.setPreviewSize(resolution[0], resolution[1])
        cam.setInterleaved(False)
        cam.setFps(fps)
        cam.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)

        xout = self.pipeline.create(dai.node.XLinkOut)
        xout.setStreamName("preview")
        cam.preview.link(xout.input)

        self.device = dai.Device(self.pipeline)
        self.q = self.device.getOutputQueue(name="preview", maxSize=4, blocking=False)
        self.mirror = mirror

    def get_frame(self):
        in_frame = self.q.tryGet()
        if in_frame is None:
            return None
        frame = in_frame.getCvFrame()
        if self.mirror:
            frame = cv2.flip(frame, 1)
        return frame

    def release(self):
        self.device.close()