# Copyright 2016 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Implements HID device interface on MacOS using IOKit and HIDManager."""

from __future__ import absolute_import

import ctypes
import ctypes.util
import logging
from six.moves.queue import Queue, Empty
import sys
import threading

from . import base

logger = logging.getLogger('_pyu2f.macos')

# Constants
DEVICE_PATH_BUFFER_SIZE = 512
DEVICE_STRING_PROPERTY_BUFFER_SIZE = 512

HID_DEVICE_PROPERTY_VENDOR_ID = b'VendorID'
HID_DEVICE_PROPERTY_PRODUCT_ID = b'ProductID'
HID_DEVICE_PROPERTY_PRODUCT = b'Product'
HID_DEVICE_PROPERTY_PRIMARY_USAGE = b'PrimaryUsage'
HID_DEVICE_PROPERTY_PRIMARY_USAGE_PAGE = b'PrimaryUsagePage'
HID_DEVICE_PROPERTY_MAX_INPUT_REPORT_SIZE = b'MaxInputReportSize'
HID_DEVICE_PROPERTY_MAX_OUTPUT_REPORT_SIZE = b'MaxOutputReportSize'
HID_DEVICE_PROPERTY_REPORT_ID = b'ReportID'


# Declare C types
class _CFType(ctypes.Structure):
  pass


class _CFString(_CFType):
  pass


class _CFSet(_CFType):
  pass


class _IOHIDManager(_CFType):
  pass


class _IOHIDDevice(_CFType):
  pass


class _CFRunLoop(_CFType):
  pass


class _CFAllocator(_CFType):
  pass


# Linter isn't consistent about valid class names. Disabling some of the errors
CF_SET_REF = ctypes.POINTER(_CFSet)
CF_STRING_REF = ctypes.POINTER(_CFString)
CF_TYPE_REF = ctypes.POINTER(_CFType)
CF_RUN_LOOP_REF = ctypes.POINTER(_CFRunLoop)
CF_RUN_LOOP_RUN_RESULT = ctypes.c_int32
CF_ALLOCATOR_REF = ctypes.POINTER(_CFAllocator)
CF_TYPE_ID = ctypes.c_ulong  # pylint: disable=invalid-name
CF_INDEX = ctypes.c_long  # pylint: disable=invalid-name
CF_TIME_INTERVAL = ctypes.c_double  # pylint: disable=invalid-name
IO_RETURN = ctypes.c_uint
IO_HID_REPORT_TYPE = ctypes.c_uint
IO_OBJECT_T = ctypes.c_uint
MACH_PORT_T = ctypes.c_uint
IO_STRING_T = ctypes.c_char_p  # pylint: disable=invalid-name
IO_SERVICE_T = IO_OBJECT_T
IO_REGISTRY_ENTRY_T = IO_OBJECT_T

IO_HID_MANAGER_REF = ctypes.POINTER(_IOHIDManager)
IO_HID_DEVICE_REF = ctypes.POINTER(_IOHIDDevice)

IO_HID_REPORT_CALLBACK = ctypes.CFUNCTYPE(None, ctypes.py_object, IO_RETURN,
                                          ctypes.c_void_p, IO_HID_REPORT_TYPE,
                                          ctypes.c_uint32,
                                          ctypes.POINTER(ctypes.c_uint8),
                                          CF_INDEX)
IO_HID_CALLBACK = ctypes.CFUNCTYPE(None, ctypes.py_object, IO_RETURN,
                                   ctypes.c_void_p)

# Define C constants
K_CF_NUMBER_SINT32_TYPE = 3
K_CF_STRING_ENCODING_UTF8 = 0x08000100
K_CF_ALLOCATOR_DEFAULT = None

K_IO_SERVICE_PLANE = b'IOService'
K_IO_MASTER_PORT_DEFAULT = 0
K_IO_HID_REPORT_TYPE_OUTPUT = 1
K_IO_RETURN_SUCCESS = 0

K_CF_RUN_LOOP_RUN_STOPPED = 2
K_CF_RUN_LOOP_RUN_TIMED_OUT = 3
K_CF_RUN_LOOP_RUN_HANDLED_SOURCE = 4

iokit = None
cf = None

# Only use iokit and cf if we're on macos, this way we can still run tests
# on other platforms if we properly mock
if sys.platform.startswith('darwin'):
  # Load relevant libraries
  iokit = ctypes.cdll.LoadLibrary(ctypes.util.find_library('IOKit'))
  cf = ctypes.cdll.LoadLibrary(ctypes.util.find_library('CoreFoundation'))

  # Exported constants
  K_CF_RUNLOOP_DEFAULT_MODE = CF_STRING_REF.in_dll(cf, 'kCFRunLoopDefaultMode')

  # Declare C function prototypes
  cf.CFSetGetValues.restype = None
  cf.CFSetGetValues.argtypes = [CF_SET_REF, ctypes.POINTER(ctypes.c_void_p)]
  cf.CFStringCreateWithCString.restype = CF_STRING_REF
  cf.CFStringCreateWithCString.argtypes = [ctypes.c_void_p, ctypes.c_char_p,
                                           ctypes.c_uint32]
  cf.CFStringGetCString.restype = ctypes.c_int
  cf.CFStringGetCString.argtypes = [CF_STRING_REF, ctypes.c_char_p, CF_INDEX,
                                    ctypes.c_uint32]
  cf.CFStringGetLength.restype = CF_INDEX
  cf.CFStringGetLength.argtypes = [CF_STRING_REF]
  cf.CFGetTypeID.restype = CF_TYPE_ID
  cf.CFGetTypeID.argtypes = [CF_TYPE_REF]
  cf.CFNumberGetTypeID.restype = CF_TYPE_ID
  cf.CFNumberGetValue.restype = ctypes.c_int
  cf.CFRelease.restype = None
  cf.CFRelease.argtypes = [CF_TYPE_REF]
  cf.CFRunLoopGetCurrent.restype = CF_RUN_LOOP_REF
  cf.CFRunLoopGetCurrent.argTypes = []
  cf.CFRunLoopRunInMode.restype = CF_RUN_LOOP_RUN_RESULT
  cf.CFRunLoopRunInMode.argtypes = [CF_STRING_REF, CF_TIME_INTERVAL,
                                    ctypes.c_bool]

  iokit.IOObjectRelease.argtypes = [IO_OBJECT_T]
  iokit.IOHIDManagerCreate.restype = IO_HID_MANAGER_REF
  iokit.IOHIDManagerCopyDevices.restype = CF_SET_REF
  iokit.IOHIDManagerCopyDevices.argtypes = [IO_HID_MANAGER_REF]
  iokit.IOHIDManagerSetDeviceMatching.restype = None
  iokit.IOHIDManagerSetDeviceMatching.argtypes = [IO_HID_MANAGER_REF,
                                                  CF_TYPE_REF]
  iokit.IOHIDDeviceGetProperty.restype = CF_TYPE_REF
  iokit.IOHIDDeviceGetProperty.argtypes = [IO_HID_DEVICE_REF, CF_STRING_REF]
  iokit.IOHIDDeviceRegisterRemovalCallback.restype = None
  iokit.IOHIDDeviceRegisterRemovalCallback.argtypes = [
      IO_HID_DEVICE_REF, IO_HID_CALLBACK, ctypes.py_object]
  iokit.IOHIDDeviceRegisterInputReportCallback.restype = None
  iokit.IOHIDDeviceRegisterInputReportCallback.argtypes = [
      IO_HID_DEVICE_REF, ctypes.POINTER(ctypes.c_uint8), CF_INDEX,
      IO_HID_REPORT_CALLBACK, ctypes.py_object]
  iokit.IORegistryEntryFromPath.restype = IO_REGISTRY_ENTRY_T
  iokit.IORegistryEntryFromPath.argtypes = [MACH_PORT_T, IO_STRING_T]
  iokit.IOHIDDeviceCreate.restype = IO_HID_DEVICE_REF
  iokit.IOHIDDeviceCreate.argtypes = [CF_ALLOCATOR_REF, IO_SERVICE_T]
  iokit.IOHIDDeviceScheduleWithRunLoop.restype = None
  iokit.IOHIDDeviceScheduleWithRunLoop.argtypes = [IO_HID_DEVICE_REF,
                                                   CF_RUN_LOOP_REF,
                                                   CF_STRING_REF]
  iokit.IOHIDDeviceUnscheduleFromRunLoop.restype = None
  iokit.IOHIDDeviceUnscheduleFromRunLoop.argtypes = [IO_HID_DEVICE_REF,
                                                     CF_RUN_LOOP_REF,
                                                     CF_STRING_REF]
  iokit.IOHIDDeviceSetReport.restype = IO_RETURN
  iokit.IOHIDDeviceSetReport.argtypes = [IO_HID_DEVICE_REF, IO_HID_REPORT_TYPE,
                                         CF_INDEX,
                                         ctypes.POINTER(ctypes.c_uint8),
                                         CF_INDEX]
else:
  logger.warn('Not running on MacOS')


def CFStr(s):
  """Builds a CFString from a python string.

  Args:
    s: source string

  Returns:
    CFStringRef representation of the source string

  Resulting CFString must be CFReleased when no longer needed.
  """
  return cf.CFStringCreateWithCString(None, s, 0)


def GetDeviceIntProperty(dev_ref, key):
  """Reads int property from the HID device."""
  cf_key = CFStr(key)
  type_ref = iokit.IOHIDDeviceGetProperty(dev_ref, cf_key)
  cf.CFRelease(cf_key)
  if not type_ref:
    return None

  if cf.CFGetTypeID(type_ref) != cf.CFNumberGetTypeID():
    raise OSError('Expected number type, got {}'.format(
        cf.CFGetTypeID(type_ref)))

  out = ctypes.c_int32()
  ret = cf.CFNumberGetValue(type_ref, K_CF_NUMBER_SINT32_TYPE,
                            ctypes.byref(out))
  if not ret:
    return None

  return out.value


def GetDeviceStringProperty(dev_ref, key):
  """Reads string property from the HID device."""
  cf_key = CFStr(key)
  type_ref = iokit.IOHIDDeviceGetProperty(dev_ref, cf_key)
  cf.CFRelease(cf_key)
  if not type_ref:
    return None

  if cf.CFGetTypeID(type_ref) != cf.CFStringGetTypeID():
    raise OSError('Expected string type, got {}'.format(
        cf.CFGetTypeID(type_ref)))

  type_ref = ctypes.cast(type_ref, CF_STRING_REF)
  out = ctypes.create_string_buffer(DEVICE_STRING_PROPERTY_BUFFER_SIZE)
  ret = cf.CFStringGetCString(type_ref, out, DEVICE_STRING_PROPERTY_BUFFER_SIZE,
                              K_CF_STRING_ENCODING_UTF8)
  if not ret:
    return None

  return out.value.decode('utf8')


def GetDevicePath(device_handle):
  """Obtains the unique path for the device.

  Args:
    device_handle: reference to the device

  Returns:
    A unique path for the device, obtained from the IO Registry

  """
  # Obtain device path from IO Registry
  io_service_obj = iokit.IOHIDDeviceGetService(device_handle)
  str_buffer = ctypes.create_string_buffer(DEVICE_PATH_BUFFER_SIZE)
  iokit.IORegistryEntryGetPath(io_service_obj, K_IO_SERVICE_PLANE, str_buffer)

  return str_buffer.value


def HidReadCallback(read_queue, result, sender, report_type, report_id, report,
                    report_length):
  """Handles incoming IN report from HID device."""
  del result, sender, report_type, report_id  # Unused by the callback function

  incoming_bytes = [report[i] for i in range(report_length)]
  read_queue.put(incoming_bytes)

# C wrapper around ReadCallback()
# Declared in this scope so it doesn't get GC-ed
REGISTERED_READ_CALLBACK = IO_HID_REPORT_CALLBACK(HidReadCallback)


def HidRemovalCallback(hid_device, result, sender):
  del result, sender
  cf.CFRunLoopStop(hid_device.run_loop_ref)


REMOVAL_CALLBACK = IO_HID_CALLBACK(HidRemovalCallback)


def DeviceReadThread(hid_device):
  """Binds a device to the thread's run loop, then starts the run loop.

  Args:
    hid_device: The MacOsHidDevice object

  The HID manager requires a run loop to handle Report reads. This thread
  function serves that purpose.
  """

  # Schedule device events with run loop
  hid_device.run_loop_ref = cf.CFRunLoopGetCurrent()
  if not hid_device.run_loop_ref:
    logger.error('Failed to get current run loop')
    return

  iokit.IOHIDDeviceScheduleWithRunLoop(hid_device.device_handle,
                                       hid_device.run_loop_ref,
                                       K_CF_RUNLOOP_DEFAULT_MODE)

  iokit.IOHIDDeviceRegisterRemovalCallback(
      hid_device.device_handle, REMOVAL_CALLBACK,
      ctypes.py_object(hid_device))

  # Run the run loop
  run_loop_run_result = cf.CFRunLoopRunInMode(
    K_CF_RUNLOOP_DEFAULT_MODE,
    4,  # Timeout in seconds
    True)  # Return after source handled

  # log any unexpected run loop exit
  if run_loop_run_result != K_CF_RUN_LOOP_RUN_HANDLED_SOURCE:
    logger.error('Unexpected run loop exit code: %d', run_loop_run_result)

  # Unschedule from run loop
  iokit.IOHIDDeviceUnscheduleFromRunLoop(hid_device.device_handle,
                                         hid_device.run_loop_ref,
                                         K_CF_RUNLOOP_DEFAULT_MODE)


class MacOsHidDevice(base.HidDevice):
  """Implementation of HID device for MacOS.

  Uses IOKit HID Manager to interact with the device.
  """

  @staticmethod
  def Enumerate():
    """See base class."""
    # Init a HID manager
    hid_mgr = iokit.IOHIDManagerCreate(None, None)
    if not hid_mgr:
      raise OSError('Unable to obtain HID manager reference')
    iokit.IOHIDManagerSetDeviceMatching(hid_mgr, None)

    # Get devices from HID manager
    device_set_ref = iokit.IOHIDManagerCopyDevices(hid_mgr)
    if not device_set_ref:
      raise OSError('Failed to obtain devices from HID manager')

    num = iokit.CFSetGetCount(device_set_ref)
    devices = (IO_HID_DEVICE_REF * num)()
    iokit.CFSetGetValues(device_set_ref, devices)

    # Retrieve and build descriptor dictionaries for each device
    descriptors = []
    for dev in devices:
      d = base.DeviceDescriptor()
      d.vendor_id = GetDeviceIntProperty(dev, HID_DEVICE_PROPERTY_VENDOR_ID)
      d.product_id = GetDeviceIntProperty(dev, HID_DEVICE_PROPERTY_PRODUCT_ID)
      d.product_string = GetDeviceStringProperty(dev,
                                                 HID_DEVICE_PROPERTY_PRODUCT)
      d.usage = GetDeviceIntProperty(dev, HID_DEVICE_PROPERTY_PRIMARY_USAGE)
      d.usage_page = GetDeviceIntProperty(
          dev, HID_DEVICE_PROPERTY_PRIMARY_USAGE_PAGE)
      d.report_id = GetDeviceIntProperty(dev, HID_DEVICE_PROPERTY_REPORT_ID)
      d.path = GetDevicePath(dev)
      descriptors.append(d.ToPublicDict())

    # Clean up CF objects
    cf.CFRelease(device_set_ref)
    cf.CFRelease(hid_mgr)

    return descriptors

  def __init__(self, path):
    # Resolve the path to device handle
    device_entry = iokit.IORegistryEntryFromPath(K_IO_MASTER_PORT_DEFAULT, path)
    if not device_entry:
      raise OSError('Device path does not match any HID device on '
                              'the system')

    self.device_handle = iokit.IOHIDDeviceCreate(K_CF_ALLOCATOR_DEFAULT,
                                                 device_entry)
    if not self.device_handle:
      raise OSError('Failed to obtain device handle from registry '
                              'entry')
    iokit.IOObjectRelease(device_entry)

    self.device_path = path

    # Open device
    result = iokit.IOHIDDeviceOpen(self.device_handle, 0)
    if result != K_IO_RETURN_SUCCESS:
      raise OSError('Failed to open device for communication: {}'
                              .format(result))

    # Create read queue
    self.read_queue = Queue()

    # Create and start read thread
    self.run_loop_ref = None

    # Read max report sizes for in/out
    self.internal_max_in_report_len = GetDeviceIntProperty(
        self.device_handle,
        HID_DEVICE_PROPERTY_MAX_INPUT_REPORT_SIZE)
    if not self.internal_max_in_report_len:
      raise OSError('Unable to obtain max in report size')

    self.internal_max_out_report_len = GetDeviceIntProperty(
        self.device_handle,
        HID_DEVICE_PROPERTY_MAX_OUTPUT_REPORT_SIZE)
    if not self.internal_max_out_report_len:
      raise OSError('Unable to obtain max out report size')

    # Register read callback
    self.in_report_buffer = (ctypes.c_uint8 * self.internal_max_in_report_len)()
    iokit.IOHIDDeviceRegisterInputReportCallback(
        self.device_handle,
        self.in_report_buffer,
        self.internal_max_in_report_len,
        REGISTERED_READ_CALLBACK,
        ctypes.py_object(self.read_queue))

  def GetInReportDataLength(self):
    """See base class."""
    return self.internal_max_in_report_len

  def GetOutReportDataLength(self):
    """See base class."""
    return self.internal_max_out_report_len

  def Write(self, packet):
    """See base class."""
    report_id = 0
    out_report_buffer = (ctypes.c_uint8 * self.internal_max_out_report_len)()
    out_report_buffer[:] = packet[:]

    result = iokit.IOHIDDeviceSetReport(self.device_handle,
                                        K_IO_HID_REPORT_TYPE_OUTPUT,
                                        report_id,
                                        out_report_buffer,
                                        self.internal_max_out_report_len)

    # Non-zero status indicates failure
    if result != K_IO_RETURN_SUCCESS:
      raise OSError('Failed to write report to device')

  def Read(self):
    """See base class."""
    read_thread = threading.Thread(target=DeviceReadThread,
                                   args=(self,))
    read_thread.start()
    read_thread.join()
    try:
        return self.read_queue.get(False)
    except Empty:
        raise OSError('Failed reading a response')

  def __del__(self):
    # Unregister the callback
    iokit.IOHIDDeviceRegisterInputReportCallback(
        self.device_handle,
        self.in_report_buffer,
        self.internal_max_in_report_len,
        ctypes.cast(0, IO_HID_REPORT_CALLBACK),
        None)
