#include "IHIDDevice.hpp"
#include <IOKit/hid/IOHIDLib.h>
#include <IOKit/hid/IOHIDDevicePlugin.h>
#include <IOKit/usb/IOUSBLib.h>
#include "IOKitPointer.hpp"
#include <thread>

namespace boo
{

class HIDDeviceIOKit : public IHIDDevice
{
    DeviceToken& m_token;
    std::shared_ptr<DeviceBase> m_devImp;

    IUnknownPointer<IOUSBInterfaceInterface> m_usbIntf;
    uint8_t m_usbIntfInPipe = 0;
    uint8_t m_usbIntfOutPipe = 0;
    CFPointer<IOHIDDeviceRef> m_hidIntf;
    bool m_runningTransferLoop = false;
    bool m_isBt = false;

    const std::string& m_devPath;
    std::mutex m_initMutex;
    std::condition_variable m_initCond;
    std::thread m_thread;

    bool _sendUSBInterruptTransfer(const uint8_t* data, size_t length)
    {
        if (m_usbIntf)
        {
            IOReturn res = m_usbIntf->WritePipe(m_usbIntf.storage(), m_usbIntfOutPipe, (void*)data, length);
            return res == kIOReturnSuccess;
        }
        return false;
    }

    size_t _receiveUSBInterruptTransfer(uint8_t* data, size_t length)
    {
        if (m_usbIntf)
        {
            UInt32 readSize = length;
            IOReturn res = m_usbIntf->ReadPipe(m_usbIntf.storage(), m_usbIntfInPipe, data, &readSize);
            if (res != kIOReturnSuccess)
                return 0;
            return readSize;
        }
        return 0;
    }

    std::vector<uint8_t> _getReportDescriptor()
    {
        if (m_hidIntf)
        {
            if (CFTypeRef desc = IOHIDDeviceGetProperty(m_hidIntf.get(), CFSTR(kIOHIDReportDescriptorKey)))
            {
                CFIndex len = CFDataGetLength(CFDataRef(desc));
                std::vector<uint8_t> ret(len, '\0');
                CFDataGetBytes(CFDataRef(desc), CFRangeMake(0, len), &ret[0]);
                return ret;
            }
        }
        return {};
    }

    bool _sendHIDReport(const uint8_t* data, size_t length, HIDReportType tp, uint32_t message)
    {
        /* HACK: A bug in IOBluetoothGamepadHIDDriver prevents raw output report transmission
         * USB driver appears to work correctly */
        if (m_hidIntf && !m_isBt)
        {
            IOReturn res = IOHIDDeviceSetReport(m_hidIntf.get(), IOHIDReportType(tp), message, data, length);
            return res == kIOReturnSuccess;
        }
        return false;
    }

    size_t _receiveHIDReport(uint8_t* data, size_t length, HIDReportType tp, uint32_t message)
    {
        if (m_hidIntf)
        {
            CFIndex readSize = length;
            IOReturn res = IOHIDDeviceGetReport(m_hidIntf.get(), IOHIDReportType(tp), message, data, &readSize);
            if (res != kIOReturnSuccess)
                return 0;
            return readSize;
        }
        return 0;
    }

    static void _threadProcUSBLL(std::shared_ptr<HIDDeviceIOKit> device)
    {
        char thrName[128];
        snprintf(thrName, 128, "%s Transfer Thread", device->m_token.getProductName().c_str());
        pthread_setname_np(thrName);
        char errStr[256];
        std::unique_lock<std::mutex> lk(device->m_initMutex);

        /* Get the HID element's parent (USB interrupt transfer-interface) */
        IOObjectPointer<io_iterator_t> devIter;
        IOObjectPointer<io_registry_entry_t> devEntry = IORegistryEntryFromPath(kIOMasterPortDefault, device->m_devPath.c_str());
        IORegistryEntryGetChildIterator(devEntry, kIOServicePlane, &devIter);
        IOObjectPointer<io_object_t> interfaceEntry;
        while (IOObjectPointer<io_service_t> obj = IOIteratorNext(devIter))
        {
            if (IOObjectConformsTo(obj, kIOUSBInterfaceClassName))
                interfaceEntry = obj;
            else
                IOObjectRelease(obj);
        }
        if (!interfaceEntry)
        {
            snprintf(errStr, 256, "Unable to find interface for %s@%s\n",
                     device->m_token.getProductName().c_str(),
                     device->m_devPath.c_str());
            device->m_devImp->deviceError(errStr);
            lk.unlock();
            device->m_initCond.notify_one();
            return;
        }

        /* IOKit Plugin COM interface (WTF Apple???) */
        IOCFPluginPointer   iodev;
        SInt32              score;
        IOReturn            err;
        err = IOCreatePlugInInterfaceForService(interfaceEntry,
                                                kIOUSBInterfaceUserClientTypeID,
                                                kIOCFPlugInInterfaceID,
                                                &iodev,
                                                &score);
        if (err)
        {
            snprintf(errStr, 256, "Unable to open %s@%s\n",
                     device->m_token.getProductName().c_str(), device->m_devPath.c_str());
            device->m_devImp->deviceError(errStr);
            lk.unlock();
            device->m_initCond.notify_one();
            return;
        }

        /* USB interface function-pointer table */
        IUnknownPointer<IOUSBInterfaceInterface> intf;
        err = iodev.As(&intf, kIOUSBInterfaceInterfaceID);
        if (err)
        {
            snprintf(errStr, 256, "Unable to open %s@%s\n",
                     device->m_token.getProductName().c_str(), device->m_devPath.c_str());
            device->m_devImp->deviceError(errStr);
            lk.unlock();
            device->m_initCond.notify_one();
            return;
        }

        /* Obtain exclusive lock on device */
        device->m_usbIntf = intf;
        err = intf->USBInterfaceOpen(intf.storage());
        if (err != kIOReturnSuccess)
        {
            if (err == kIOReturnExclusiveAccess)
            {
                snprintf(errStr, 256, "Unable to open %s@%s: someone else using it\n",
                         device->m_token.getProductName().c_str(), device->m_devPath.c_str());
                device->m_devImp->deviceError(errStr);
            }
            else
            {
                snprintf(errStr, 256, "Unable to open %s@%s\n",
                         device->m_token.getProductName().c_str(), device->m_devPath.c_str());
                device->m_devImp->deviceError(errStr);
            }
            lk.unlock();
            device->m_initCond.notify_one();
            return;
        }

        /* Determine pipe indices for interrupt I/O */
        UInt8 numEndpoints = 0;
        err = intf->GetNumEndpoints(intf.storage(), &numEndpoints);
        for (int i=1 ; i<numEndpoints+1 ; ++i)
        {
            UInt8 dir, num, tType, interval;
            UInt16 mPacketSz;
            err = intf->GetPipeProperties(intf.storage(), i, &dir, &num, &tType, &mPacketSz, &interval);
            if (tType == kUSBInterrupt)
            {
                if (dir == kUSBIn)
                    device->m_usbIntfInPipe = num;
                else if (dir == kUSBOut)
                    device->m_usbIntfOutPipe = num;
            }
        }

        /* Return control to main thread */
        device->m_runningTransferLoop = true;
        lk.unlock();
        device->m_initCond.notify_one();

        /* Start transfer loop */
        device->m_devImp->initialCycle();
        while (device->m_runningTransferLoop)
            device->m_devImp->transferCycle();
        device->m_devImp->finalCycle();

        /* Cleanup */
        intf->USBInterfaceClose(intf.storage());
        device->m_usbIntf = nullptr;
    }

    static void _threadProcBTLL(std::shared_ptr<HIDDeviceIOKit> device)
    {
        std::unique_lock<std::mutex> lk(device->m_initMutex);

        /* Return control to main thread */
        device->m_runningTransferLoop = true;
        lk.unlock();
        device->m_initCond.notify_one();

        /* Start transfer loop */
        device->m_devImp->initialCycle();
        while (device->m_runningTransferLoop)
            device->m_devImp->transferCycle();
        device->m_devImp->finalCycle();
    }

    static void _hidRemoveCb(void * _Nullable        context,
                             IOReturn                result,
                             void * _Nullable        sender)
    {
        reinterpret_cast<HIDDeviceIOKit*>(context)->m_runningTransferLoop = false;
    }

    static void _hidReportCb(void * _Nullable        context,
                             IOReturn,
                             void * _Nullable,
                             IOHIDReportType         type,
                             uint32_t                reportID,
                             uint8_t *               report,
                             CFIndex                 reportLength)
    {
        reinterpret_cast<DeviceBase*>(context)->receivedHIDReport(report, reportLength, HIDReportType(type), reportID);
    }

    static void _threadProcHID(std::shared_ptr<HIDDeviceIOKit> device)
    {
        char thrName[128];
        snprintf(thrName, 128, "%s Transfer Thread", device->m_token.getProductName().c_str());
        pthread_setname_np(thrName);
        char errStr[256];
        std::unique_lock<std::mutex> lk(device->m_initMutex);

        /* Get the HID element's object (HID device interface) */
        IOObjectPointer<io_service_t> interfaceEntry =
            IORegistryEntryFromPath(kIOMasterPortDefault, device->m_devPath.c_str());
        if (!IOObjectConformsTo(interfaceEntry.get(), "IOHIDDevice"))
        {
            snprintf(errStr, 256, "Unable to find interface for %s@%s\n",
                     device->m_token.getProductName().c_str(),
                     device->m_devPath.c_str());
            device->m_devImp->deviceError(errStr);
            lk.unlock();
            device->m_initCond.notify_one();
            return;
        }

        device->m_hidIntf = IOHIDDeviceCreate(nullptr, interfaceEntry.get());
        if (!device->m_hidIntf)
        {
            snprintf(errStr, 256, "Unable to open %s@%s\n",
                     device->m_token.getProductName().c_str(), device->m_devPath.c_str());
            device->m_devImp->deviceError(errStr);
            lk.unlock();
            device->m_initCond.notify_one();
            return;
        }

        /* Open device */
        IOReturn err = IOHIDDeviceOpen(device->m_hidIntf.get(), kIOHIDOptionsTypeNone);
        if (err != kIOReturnSuccess)
        {
            if (err == kIOReturnExclusiveAccess)
            {
                snprintf(errStr, 256, "Unable to open %s@%s: someone else using it\n",
                         device->m_token.getProductName().c_str(), device->m_devPath.c_str());
                device->m_devImp->deviceError(errStr);
            }
            else
            {
                snprintf(errStr, 256, "Unable to open %s@%s\n",
                         device->m_token.getProductName().c_str(), device->m_devPath.c_str());
                device->m_devImp->deviceError(errStr);
            }
            lk.unlock();
            device->m_initCond.notify_one();
            return;
        }

        /* Register removal callback */
        IOHIDDeviceRegisterRemovalCallback(device->m_hidIntf.get(), _hidRemoveCb, device.get());

        /* Make note if device uses bluetooth driver */
        if (CFTypeRef transport = IOHIDDeviceGetProperty(device->m_hidIntf.get(), CFSTR(kIOHIDTransportKey)))
            device->m_isBt = CFStringCompare(CFStringRef(transport), CFSTR(kIOHIDTransportBluetoothValue), 0) == kCFCompareEqualTo;

        /* Register input buffer */
        std::unique_ptr<uint8_t[]> buffer;
        int bufSize = 0;
        if (CFTypeRef maxSize = IOHIDDeviceGetProperty(device->m_hidIntf.get(), CFSTR(kIOHIDMaxInputReportSizeKey)))
            CFNumberGetValue(CFNumberRef(maxSize), kCFNumberIntType, &bufSize);
        if (bufSize)
        {
            buffer = std::unique_ptr<uint8_t[]>(new uint8_t[bufSize]);
            IOHIDDeviceRegisterInputReportCallback(device->m_hidIntf.get(), buffer.get(), bufSize,
                                                   _hidReportCb, device->m_devImp.get());
            IOHIDDeviceScheduleWithRunLoop(device->m_hidIntf.get(), CFRunLoopGetCurrent(), kCFRunLoopDefaultMode);
        }

        /* Return control to main thread */
        device->m_runningTransferLoop = true;
        lk.unlock();
        device->m_initCond.notify_one();

        /* Start transfer loop */
        device->m_devImp->initialCycle();
        while (device->m_runningTransferLoop)
        {
            CFRunLoopRunInMode(kCFRunLoopDefaultMode, 0.010, true);
            if (device->m_runningTransferLoop)
                device->m_devImp->transferCycle();
        }
        device->m_devImp->finalCycle();

        /* Cleanup */
        IOHIDDeviceClose(device->m_hidIntf.get(), kIOHIDOptionsTypeNone);
        device->m_hidIntf.reset();
    }

    void _deviceDisconnected()
    {
        m_runningTransferLoop = false;
    }

public:

    HIDDeviceIOKit(DeviceToken& token, const std::shared_ptr<DeviceBase>& devImp)
    : m_token(token),
      m_devImp(devImp),
      m_devPath(token.getDevicePath())
    {
    }

    void _startThread()
    {
        std::unique_lock<std::mutex> lk(m_initMutex);
        DeviceType dType = m_token.getDeviceType();
        if (dType == DeviceType::USB)
            m_thread = std::thread(_threadProcUSBLL, std::static_pointer_cast<HIDDeviceIOKit>(shared_from_this()));
        else if (dType == DeviceType::Bluetooth)
            m_thread = std::thread(_threadProcBTLL, std::static_pointer_cast<HIDDeviceIOKit>(shared_from_this()));
        else if (dType == DeviceType::HID)
            m_thread = std::thread(_threadProcHID, std::static_pointer_cast<HIDDeviceIOKit>(shared_from_this()));
        else
        {
            fprintf(stderr, "invalid token supplied to device constructor\n");
            return;
        }
        m_initCond.wait(lk);
    }

    ~HIDDeviceIOKit()
    {
        m_runningTransferLoop = false;
        if (m_thread.joinable())
            m_thread.detach();
    }
};

std::shared_ptr<IHIDDevice> IHIDDeviceNew(DeviceToken& token, const std::shared_ptr<DeviceBase>& devImp)
{
    return std::make_shared<HIDDeviceIOKit>(token, devImp);
}

}