import cv2


class NVGstreamer:

    def __init__(self, buildconfig):
        self.width = 480
        self.height = 640
        self.latency = 0
        self.framerate = "10/1"
        self.fformat = "BGRx"

        self.BUILD_CONFIG = {
            "width": self.width,
            "height": self.height,
            "latency": self.latency,
            "framerate": self.framerate,
            "format": self.fformat,
            "gstreamer": True
        }

        self.BUILD_CONFIG.update(buildconfig)

    def open_cam_rtsp(self):
        gst_str = ('rtspsrc location={uri} latency={latency} ! '
                   'rtph264depay ! h264parse ! omxh264dec ! '
                   'nvvidconv ! videorate ! '
                   'video/x-raw, width=(int){width}, height=(int){height}, '
                   'format=(string){format}, framerate=(fraction){framerate} ! '
                   'videoconvert ! appsink').format(**self.BUILD_CONFIG)
        print(gst_str)
        return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)

    def open_cam_usb(self):
        # We want to set width and height here, otherwise we could just do:
        #     return cv2.VideoCapture(dev)
        gst_str = ('v4l2src device=/dev/video{uri} ! '
                   'video/x-raw, width=(int){width}, height=(int){height}, '
                   'format=(string){format}, framerate=(fraction){framerate} ! '
                   'videoconvert ! appsink').format(**self.BUILD_CONFIG)
        print(gst_str)
        return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)

    def open_cam_onboard(self):
        # On versions of L4T prior to 28.1, add 'flip-method=2' into gst_str
        gst_str = ('nvcamerasrc ! '
                   'video/x-raw(memory:NVMM), '
                   'width=(int)2592, height=(int)1458, '
                   'format=(string)I420 ! '
                   'nvvidconv ! videorate ! '
                   'video/x-raw, width=(int){width}, height=(int){height}, '
                   'format=(string){format}, framerate=(fraction){framerate} !'
                   'videoconvert ! appsink').format(**self.BUILD_CONFIG)
        print(gst_str)
        return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)

    def custom_pipeline(self):
        gst_str = "{customGstPipelineString}".format(**self.BUILD_CONFIG)
        print(gst_str)
        return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)

    def build_pipeline(self):
        if self.BUILD_CONFIG["gStreamer"]!=True:
            if  self.BUILD_CONFIG["sourceType"] == "usbcam":
                self.cap = cv2.VideoCapture(int(self.BUILD_CONFIG["uri"]))
            else:
                self.cap = cv2.VideoCapture(self.BUILD_CONFIG["uri"])
        elif self.BUILD_CONFIG["sourceType"] == "rtsp":
            self.cap = self.open_cam_rtsp()
        elif self.BUILD_CONFIG["sourceType"] == "usbcam":
            self.cap = self.open_cam_usb()
        elif self.BUILD_CONFIG["sourceType"] == "onboard":
            self.cap = self.open_cam_onboard()
        elif self.BUILD_CONFIG["sourceType"] == "customPipeline":
            self.cap = self.custom_pipeline()
        else:
            raise ValueError("unimplemented source {}".format(self.BUILD_CONFIG["sourceType"]))

    def get_stream(self):
        return self.cap
