#!/usr/bin/env python3

import depthai as dai
import av
import cv2
# Create pipeline
pipeline = dai.Pipeline()

camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setBoardSocket(dai.CameraBoardSocket.RGB)
camRgb.setFps(30)
camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)

def create_encoder(name):
    videoEnc = pipeline.create(dai.node.VideoEncoder)
    videoEnc.setDefaultProfilePreset(camRgb.getFps(), dai.VideoEncoderProperties.Profile.H264_MAIN)
    camRgb.video.link(videoEnc.input)

    xout = pipeline.create(dai.node.XLinkOut)
    xout.setStreamName(name)
    videoEnc.bitstream.link(xout.input)

create_encoder("1")
create_encoder("2")
create_encoder("3")
create_encoder("4")

from depthai_sdk.fps import FPSHandler
fps = FPSHandler()

with dai.Device(pipeline) as device:
    qs = [device.getOutputQueue(name, maxSize=30, blocking=True) for name in ["1", "2", "3", "4"]]
    codec = av.CodecContext.create("h264", "r")
    stop = False
    while not stop:
        for i, q in enumerate(qs):
            data = q.get().getData()  # Blocking call, will wait until new data has arrived

            if i == 0:
                fps.nextIter()
                print(fps.fps())
                packets = codec.parse(data)

                for packet in packets:
                    frames = codec.decode(packet)
                    if frames:
                        frame = frames[0].to_ndarray(format='bgr24')
                        cv2.imshow('frame', frame)

                if cv2.waitKey(1) == ord('q'):
                    stop = True