#!/usr/bin/env python3

import cv2
import depthai as dai
import contextlib

# Start defining a pipeline
print("Define pipeline")
pipeline = dai.Pipeline()

# Define a source - color camera
print("Define a source")
cam_rgb = pipeline.createColorCamera()
cam_rgb.setPreviewSize(600, 600)
cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)
#cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_4_K)
cam_rgb.setInterleaved(False)
video_enc = pipeline.createVideoEncoder()

# Properties
print("Define properties")
video_enc.setDefaultProfilePreset(3840, 2160, 30, dai.VideoEncoderProperties.Profile.H265_MAIN)

# Create output
print("Create output")
xout_rgb = pipeline.createXLinkOut()
xout_rgb.setStreamName("h265")
cam_rgb.preview.link(xout_rgb.input)
cam_rgb.video.link(video_enc.input)
video_enc.bitstream.link(xout_rgb.input)

q_rgb_list = []
print("q_rgb_list")

# https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack
with contextlib.ExitStack() as stack:
    for device_info in dai.Device.getAllAvailableDevices():
        print("for device_info...")
        device = stack.enter_context(dai.Device(pipeline, device_info))
        deviceIP = device_info.getMxId()
        print("Device IP: " + deviceIP)
        print("Conected to " + device_info.getMxId())
        device.startPipeline()
        # Output queue will be used to get the rgb frames from the output defined above
        #q_rgb = device.getOutputQueue(name="h265", maxSize=4, blocking=False)
        # Output queue will be used to get the encoded data from the output defined above
        #q_rgb = device.getOutputQueue(name="h265", maxSize=30, blocking=True)
        q_rgb = device.getOutputQueue(name="h265", maxSize=30, blocking=False)
        #open('video.h265.' + deviceIP, 'wb') as videoFile
#        open('video.h265', 'wb') as videoFile
        q_rgb_list.append(q_rgb)
        #with open('video.h265', 'wb') as videoFile:
        open('video.h265', 'wb') as videoFile
            while True:
                h265Packet = q.get()  # Blocking call, will wait until a new data has arrived
                h265Packet.getData().tofile(videoFile)  # Appends the packet data to the opened file

    while True:
        for i, q_rgb in enumerate(q_rgb_list):
            print("q_rgb in enumerate...")
            in_rgb = q_rgb.tryGet()
            if in_rgb is not None:
#                print("imshow...")
#                cv2.imshow("rgb-" + str(i + 1), in_rgb.getCvFrame())
                print("get packet...")
                h265Packet = q.get()  # Blocking call, will wait until a new data has arrived
                print("append packet...")
                h265Packet.getData().tofile(videoFile)  # Appends the packet data to the opened file
        #with open('video.h265', 'wb') as videoFile:
        #    try:
        #        while True:
        #            h265Packet = q.get()  # Blocking call, will wait until a new data has arrived
        #            h265Packet.getData().tofile(videoFile)  # Appends the packet data to the opened file

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

###### XXX ENCODING XXX ######
#with dai.Device(pipeline) as device:
#    with open('video.h265', 'wb') as videoFile:
#        print("Press Ctrl+C to stop encoding...")
#        try:
#            while True:
#                h265Packet = q.get()  # Blocking call, will wait until a new data has arrived
#                h265Packet.getData().tofile(videoFile)  # Appends the packet data to the opened file