import cv2
import numpy

import urllib


class MJPEGGrabber(object):
    SOI = '\xff\xd8'
    EOI = '\r\n--'

    CHUNKSIZE = 10 * 1024  # 10 KiB

    def __init__(self, conn):
        self.conn = conn

        self._buf = bytearray()
        self._eoi = -1

    def update(self):
        self._buf += self.conn.read(self.CHUNKSIZE)

        self._eoi = self._buf.rfind(self.EOI)

        soi = self._buf.rfind(self.SOI, 0, self._eoi)

        if soi != -1:
            self._buf = self._buf[soi:]
        else:
            self._eoi = -1

        return True

    def get_image(self):
        # make sure we have an image in buffer first
        while self._eoi == -1:
            self.update()

        jpeg = str(self._buf[:self._eoi])
        jpeg = numpy.fromstring(jpeg, dtype=numpy.uint8)
        jpeg = cv2.imdecode(jpeg, cv2.IMREAD_COLOR)

        self._buf = self._buf[self._eoi+4:]
        self._eoi = -1

        return jpeg

    def close(self):
        del self._buf


class IPCamGrabber(MJPEGGrabber):
    def __init__(self, url):
        conn = urllib.urlopen(url)
        super(IPCamGrabber, self).__init__(conn)

    def close(self):
        self.conn.close()
        super(IPCamGrabber, self).close()

if __name__ == '__main__':
    import sys

    def usage():
        print 'Usage: {} [url]'.format(sys.argv[0])
        sys.exit(1)

    sources = sys.argv[1:]
    if len(sources) == 0:
        sources = ['http://robot:8080?action=stream']

    streams = map(IPCamGrabber, sources)

    print 'Press any key to exit...'

    try:
        while 1:
            for name, mjpeg in zip(sources, streams):
                im = mjpeg.get_image()
                cv2.imshow(name, im)
            if cv2.waitKey(5) != -1:
                break
    except KeyboardInterrupt:
        pass
    finally:
        # mjpeg.close()
        cv2.destroyAllWindows()
