#!/usr/bin/env python3
#
# This file is part of Checkbox.
#
# Copyright 2008-2012 Canonical Ltd.
#
# The v4l2 ioctl code comes from the Python bindings for the v4l2
# userspace api (http://pypi.python.org/pypi/v4l2):
# Copyright (C) 1999-2009 the contributors
#
# The JPEG metadata parser is a part of bfg-pages:
# http://code.google.com/p/bfg-pages/source/browse/trunk/pages/getimageinfo.py
# Copyright (C) Tim Hoffman
#
# Checkbox is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3,
# as published by the Free Software Foundation.

#
# Checkbox is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Checkbox.  If not, see <http://www.gnu.org/licenses/>.
#

import os
import re
import sys
import time
import errno
import fcntl
import ctypes
import struct
import imghdr
from tempfile import NamedTemporaryFile
from subprocess import check_call, CalledProcessError, STDOUT
import argparse
from glob import glob
from gi.repository import GObject


_IOC_NRBITS = 8
_IOC_TYPEBITS = 8
_IOC_SIZEBITS = 14

_IOC_NRSHIFT = 0
_IOC_TYPESHIFT = _IOC_NRSHIFT + _IOC_NRBITS
_IOC_SIZESHIFT = _IOC_TYPESHIFT + _IOC_TYPEBITS
_IOC_DIRSHIFT = _IOC_SIZESHIFT + _IOC_SIZEBITS

_IOC_WRITE = 1
_IOC_READ = 2


def _IOC(dir_, type_, nr, size):
    return (
        ctypes.c_int32(dir_ << _IOC_DIRSHIFT).value |
        ctypes.c_int32(ord(type_) << _IOC_TYPESHIFT).value |
        ctypes.c_int32(nr << _IOC_NRSHIFT).value |
        ctypes.c_int32(size << _IOC_SIZESHIFT).value)


def _IOC_TYPECHECK(t):
    return ctypes.sizeof(t)


def _IOR(type_, nr, size):
    return _IOC(_IOC_READ, type_, nr, ctypes.sizeof(size))


def _IOWR(type_, nr, size):
    return _IOC(_IOC_READ | _IOC_WRITE, type_, nr, _IOC_TYPECHECK(size))


class v4l2_capability(ctypes.Structure):
    """
    Driver capabilities
    """
    _fields_ = [
        ('driver', ctypes.c_char * 16),
        ('card', ctypes.c_char * 32),
        ('bus_info', ctypes.c_char * 32),
        ('version', ctypes.c_uint32),
        ('capabilities', ctypes.c_uint32),
        ('reserved', ctypes.c_uint32 * 4),
    ]


# Values for 'capabilities' field
V4L2_CAP_VIDEO_CAPTURE = 0x00000001
V4L2_CAP_VIDEO_OVERLAY = 0x00000004
V4L2_CAP_READWRITE = 0x01000000
V4L2_CAP_STREAMING = 0x04000000

v4l2_frmsizetypes = ctypes.c_uint
(
    V4L2_FRMSIZE_TYPE_DISCRETE,
    V4L2_FRMSIZE_TYPE_CONTINUOUS,
    V4L2_FRMSIZE_TYPE_STEPWISE,
) = range(1, 4)


class v4l2_frmsize_discrete(ctypes.Structure):
    _fields_ = [
        ('width', ctypes.c_uint32),
        ('height', ctypes.c_uint32),
    ]


class v4l2_frmsize_stepwise(ctypes.Structure):
    _fields_ = [
        ('min_width', ctypes.c_uint32),
        ('min_height', ctypes.c_uint32),
        ('step_width', ctypes.c_uint32),
        ('min_height', ctypes.c_uint32),
        ('max_height', ctypes.c_uint32),
        ('step_height', ctypes.c_uint32),
    ]


class v4l2_frmsizeenum(ctypes.Structure):
    class _u(ctypes.Union):
        _fields_ = [
            ('discrete', v4l2_frmsize_discrete),
            ('stepwise', v4l2_frmsize_stepwise),
        ]

    _fields_ = [
        ('index', ctypes.c_uint32),
        ('pixel_format', ctypes.c_uint32),
        ('type', ctypes.c_uint32),
        ('_u', _u),
        ('reserved', ctypes.c_uint32 * 2)
    ]

    _anonymous_ = ('_u',)


class v4l2_fmtdesc(ctypes.Structure):
    _fields_ = [
        ('index', ctypes.c_uint32),
        ('type', ctypes.c_int),
        ('flags', ctypes.c_uint32),
        ('description', ctypes.c_char * 32),
        ('pixelformat', ctypes.c_uint32),
        ('reserved', ctypes.c_uint32 * 4),
    ]

V4L2_FMT_FLAG_COMPRESSED = 0x0001
V4L2_FMT_FLAG_EMULATED = 0x0002

# ioctl code for video devices
VIDIOC_QUERYCAP = _IOR('V', 0, v4l2_capability)
VIDIOC_ENUM_FRAMESIZES = _IOWR('V', 74, v4l2_frmsizeenum)
VIDIOC_ENUM_FMT = _IOWR('V', 2, v4l2_fmtdesc)


class CameraTest:
    """
    A simple class that displays a test image via GStreamer.
    """
    def __init__(self, args, gst_plugin=None, gst_video_type=None):
        self.args = args
        self._mainloop = GObject.MainLoop()
        self._width = 640
        self._height = 480
        self._gst_plugin = gst_plugin
        self._gst_video_type = gst_video_type

    def detect(self):
        """
        Display information regarding webcam hardware
        """
        cap_status = dev_status = 1
        for i in range(10):
            cp = v4l2_capability()
            device = '/dev/video%d' % i
            try:
                with open(device, 'r') as vd:
                    fcntl.ioctl(vd, VIDIOC_QUERYCAP, cp)
            except IOError:
                continue
            dev_status = 0
            print("%s: OK" % device)
            print("    name   : %s" % cp.card.decode('UTF-8'))
            print("    driver : %s" % cp.driver.decode('UTF-8'))
            print("    version: %s.%s.%s"
                  % (cp.version >> 16,
                  (cp.version >> 8) & 0xff,
                  cp.version & 0xff))
            print("    flags  : 0x%x [" % cp.capabilities,
                  ' CAPTURE' if cp.capabilities & V4L2_CAP_VIDEO_CAPTURE
                  else '',
                  ' OVERLAY' if cp.capabilities & V4L2_CAP_VIDEO_OVERLAY
                  else '',
                  ' READWRITE' if cp.capabilities & V4L2_CAP_READWRITE
                  else '',
                  ' STREAMING' if cp.capabilities & V4L2_CAP_STREAMING
                  else '',
                  ' ]', sep="")

            resolutions = self._get_supported_resolutions(device)
            print('    ',
                  self._supported_resolutions_to_string(resolutions).replace(
                  "\n", " "),
                  sep="")

            if cp.capabilities & V4L2_CAP_VIDEO_CAPTURE:
                cap_status = 0
        return dev_status | cap_status

    def led(self):
        """
        Activate camera (switch on led), but don't display any output
        """
        pipespec = ("v4l2src device=%(device)s "
                    "! %(type)s "
                    "! %(plugin)s "
                    "! testsink"
                    % {'device': self.args.device,
                       'type': self._gst_video_type,
                       'plugin': self._gst_plugin})
        self._pipeline = Gst.parse_launch(pipespec)
        self._pipeline.set_state(Gst.State.PLAYING)
        time.sleep(10)
        self._pipeline.set_state(Gst.State.NULL)

    def display(self):
        """
        Displays the preview window
        """
        pipespec = ("v4l2src device=%(device)s "
                    "! %(type)s,width=%(width)d,height=%(height)d "
                    "! %(plugin)s "
                    "! autovideosink"
                    % {'device': self.args.device,
                       'type': self._gst_video_type,
                       'width': self._width,
                       'height': self._height,
                       'plugin': self._gst_plugin})
        self._pipeline = Gst.parse_launch(pipespec)
        self._pipeline.set_state(Gst.State.PLAYING)
        time.sleep(10)
        self._pipeline.set_state(Gst.State.NULL)

    def still(self):
        """
        Captures an image to a file
        """
        if self.args.filename:
            self._still_helper(self.args.filename, self._width, self._height,
                               self.args.quiet)
        else:
            with NamedTemporaryFile(prefix='camera_test_', suffix='.jpg') as f:
                self._still_helper(f.name, self._width, self._height,
                                   self.args.quiet)

    def _still_helper(self, filename, width, height, quiet, pixelformat=None):
        """
        Captures an image to a given filename.  width and height specify the
        image size and quiet controls whether the image is displayed to the
        user (quiet = True means do not display image).
        """
        command = ["fswebcam", "-D 1", "-S 50", "--no-banner",
                   "-d", self.args.device,
                   "-r", "%dx%d"
                   % (width, height), filename]
        use_gstreamer = False
        if pixelformat:
            command.extend(["-p", pixelformat])

        try:
            check_call(command, stdout=open(os.devnull, 'w'), stderr=STDOUT)
        except (CalledProcessError, OSError):
            use_gstreamer = True

        if use_gstreamer:
            pipespec = ("v4l2src device=%(device)s "
                        "! %(type)s,width=%(width)d,height=%(height)d "
                        "! %(plugin)s "
                        "! jpegenc "
                        "! filesink location=%(filename)s"
                        % {'device': self.args.device,
                           'type': self._gst_video_type,
                           'width': width,
                           'height': height,
                           'plugin': self._gst_plugin,
                           'filename': filename})
            self._pipeline = Gst.parse_launch(pipespec)
            self._pipeline.set_state(Gst.State.PLAYING)
            time.sleep(3)
            self._pipeline.set_state(Gst.State.NULL)

        if not quiet:
            try:
                check_call(["timeout", "-k", "11", "10", "eog", filename])
            except CalledProcessError:
                pass

    def _supported_resolutions_to_string(self, supported_resolutions):
        """
        Return a printable string representing a list of supported resolutions
        """
        ret = ""
        for resolution in supported_resolutions:
            ret += "Format: %s (%s)\n" % (resolution['pixelformat'],
                   resolution['description'])
            ret += "Resolutions: "
            for res in resolution['resolutions']:
                ret += "%sx%s," % (res[0], res[1])
            # truncate the extra comma with :-1
            ret = ret[:-1] + "\n"
        return ret

    def resolutions(self):
        """
        After querying the webcam for supported formats and resolutions,
        take multiple images using the first format returned by the driver,
        and see if they are valid
        """
        resolutions = self._get_supported_resolutions(self.args.device)
        # print supported formats and resolutions for the logs
        print(self._supported_resolutions_to_string(resolutions))

        # pick the first format, which seems to be what the driver wants for a
        # default.  This also matches the logic that fswebcam uses to select
        # a default format.
        resolution = resolutions[0]
        if resolution:
            print("Taking multiple images using the %s format"
                  % resolution['pixelformat'])
            for res in resolution['resolutions']:
                w = res[0]
                h = res[1]
                f = NamedTemporaryFile(prefix='camera_test_%s%sx%s' %
                                       (resolution['pixelformat'], w, h),
                                       suffix='.jpg', delete=False)
                print("Taking a picture at %sx%s" % (w, h))
                self._still_helper(f.name, w, h, True,
                                   pixelformat=resolution['pixelformat'])
                if self._validate_image(f.name, w, h):
                    print("Validated image %s" % f.name)
                    os.remove(f.name)
                else:
                    print("Failed to validate image %s" % f.name,
                          file=sys.stderr)
                    os.remove(f.name)
                    return 1
            return 0

    def _get_pixel_formats(self, device, maxformats=5):
        """
        Query the camera to see what pixel formats it supports.  A list of
        dicts is returned consisting of format and description.  The caller
        should check whether this camera supports VIDEO_CAPTURE before
        calling this function.
        """
        supported_formats = []
        fmt = v4l2_fmtdesc()
        fmt.index = 0
        fmt.type = V4L2_CAP_VIDEO_CAPTURE
        try:
            while fmt.index < maxformats:
                with open(device, 'r') as vd:
                    if fcntl.ioctl(vd, VIDIOC_ENUM_FMT, fmt) == 0:
                        pixelformat = {}
                        # save the int type for re-use later
                        pixelformat['pixelformat_int'] = fmt.pixelformat
                        pixelformat['pixelformat'] = "%s%s%s%s" % \
                            (chr(fmt.pixelformat & 0xFF),
                             chr((fmt.pixelformat >> 8) & 0xFF),
                             chr((fmt.pixelformat >> 16) & 0xFF),
                             chr((fmt.pixelformat >> 24) & 0xFF))
                        pixelformat['description'] = fmt.description.decode()
                        supported_formats.append(pixelformat)
                fmt.index = fmt.index + 1
        except IOError as e:
            # EINVAL is the ioctl's way of telling us that there are no
            # more formats, so we ignore it
            if e.errno != errno.EINVAL:
                print("Unable to determine Pixel Formats, this may be a "
                      "driver issue.")
            return supported_formats
        return supported_formats

    def _get_supported_resolutions(self, device):
        """
        Query the camera for supported resolutions for a given pixel_format.
        Data is returned in a list of dictionaries with supported pixel
        formats as the following example shows:
        resolution['pixelformat'] = "YUYV"
        resolution['description'] = "(YUV 4:2:2 (YUYV))"
        resolution['resolutions'] = [[width, height], [640, 480], [1280, 720] ]

        If we are unable to gather any information from the driver, then we
        return YUYV and 640x480 which seems to be a safe default.
        Per the v4l2 spec the ioctl used here is experimental
        but seems to be well supported.
        """
        supported_formats = self._get_pixel_formats(device)
        if not supported_formats:
            resolution = {}
            resolution['description'] = "YUYV"
            resolution['pixelformat'] = "YUYV"
            resolution['resolutions'] = [[640, 480]]
            supported_formats.append(resolution)
            return supported_formats

        for supported_format in supported_formats:
            resolutions = []
            framesize = v4l2_frmsizeenum()
            framesize.index = 0
            framesize.pixel_format = supported_format['pixelformat_int']
            with open(device, 'r') as vd:
                try:
                    while fcntl.ioctl(vd,
                                      VIDIOC_ENUM_FRAMESIZES,
                                      framesize) == 0:
                        if framesize.type == V4L2_FRMSIZE_TYPE_DISCRETE:
                            resolutions.append([framesize.discrete.width,
                                               framesize.discrete.height])
                        # for continuous and stepwise, let's just use min and
                        # max they use the same structure and only return
                        # one result
                        elif framesize.type == V4L2_FRMSIZE_TYPE_CONTINUOUS or\
                            framesize.type == V4L2_FRMSIZE_TYPE_STEPWISE:
                            resolutions.append([framesize.stepwise.min_width,
                                                framesize.stepwise.min_height]
                                               )
                            resolutions.append([framesize.stepwise.max_width,
                                                framesize.stepwise.max_height]
                                               )
                            break
                        framesize.index = framesize.index + 1
                except IOError as e:
                    # EINVAL is the ioctl's way of telling us that there are no
                    # more formats, so we ignore it
                    if e.errno != errno.EINVAL:
                        print("Unable to determine supported framesizes "
                              "(resolutions), this may be a driver issue.")
            supported_format['resolutions'] = resolutions
        return supported_formats

    def _validate_image(self, filename, width, height):
        """
        Given a filename, ensure that the image is the width and height
        specified and is a valid image file.
        """
        if imghdr.what(filename) != 'jpeg':
            return False

        outw = outh = 0
        with open(filename, mode='rb') as jpeg:
            jpeg.seek(2)
            b = jpeg.read(1)
            try:
                while (b and ord(b) != 0xDA):
                    while (ord(b) != 0xFF):
                        b = jpeg.read(1)
                    while (ord(b) == 0xFF):
                        b = jpeg.read(1)
                    if (ord(b) >= 0xC0 and ord(b) <= 0xC3):
                        jpeg.seek(3, 1)
                        h, w = struct.unpack(">HH", jpeg.read(4))
                        break
                    b = jpeg.read(1)
                outw, outh = int(w), int(h)
            except (struct.error, ValueError):
                pass

            if outw != width:
                print("Image width does not match, was %s should be %s" %
                      (outw, width), file=sys.stderr)
                return False
            if outh != height:
                print("Image width does not match, was %s should be %s" %
                      (outh, height), file=sys.stderr)
                return False

            return True

        return True


def parse_arguments(argv):
    """
    Parse command line arguments
    """
    parser = argparse.ArgumentParser(description="Run a camera-related test")
    subparsers = parser.add_subparsers(dest='test',
                                       title='test',
                                       description='Available camera tests')

    def add_device_parameter(parser):
        group = parser.add_mutually_exclusive_group()
        group.add_argument("-d", "--device", default="/dev/video0",
                           help="Device for the webcam to use")
        group.add_argument("--highest-device", action="store_true",
                           help=("Use the /dev/videoN "
                                 "where N is the highest value available"))
        group.add_argument("--lowest-device", action="store_true",
                           help=("Use the /dev/videoN "
                                 "where N is the lowest value available"))

    subparsers.add_parser('detect')
    led_parser = subparsers.add_parser('led')
    add_device_parameter(led_parser)
    display_parser = subparsers.add_parser('display')
    add_device_parameter(display_parser)
    still_parser = subparsers.add_parser('still')
    add_device_parameter(still_parser)
    still_parser.add_argument("-f", "--filename",
                              help="Filename to store the picture")
    still_parser.add_argument("-q", "--quiet", action="store_true",
                              help=("Don't display picture, "
                                    "just write the picture to a file"))
    resolutions_parser = subparsers.add_parser('resolutions')
    add_device_parameter(resolutions_parser)
    args = parser.parse_args(argv)

    def get_video_devices():
        devices = sorted(glob('/dev/video[0-9]'),
                         key=lambda d: re.search(r'\d', d).group(0))
        assert len(devices) > 0, "No video devices found"
        return devices

    if hasattr(args, 'highest_device') and args.highest_device:
        args.device = get_video_devices()[-1]
    elif hasattr(args, 'lowest_device') and args.lowest_device:
        args.device = get_video_devices()[0]
    return args


if __name__ == "__main__":
    args = parse_arguments(sys.argv[1:])

    if not args.test:
        args.test = 'detect'

    # Import Gst only for the test cases that will need it
    if args.test in ['display', 'still', 'led', 'resolutions']:
        from gi.repository import Gst
        if Gst.version()[0] > 0:
            gst_plugin = 'videoconvert'
            gst_video_type = 'video/x-raw'
        else:
            gst_plugin = 'ffmpegcolorspace'
            gst_video_type = 'video/x-raw-yuv'
        Gst.init(None)
        camera = CameraTest(args, gst_plugin, gst_video_type)
    else:
        camera = CameraTest(args)

    sys.exit(getattr(camera, args.test)())
