// USB transport (IOKit / IOUSBFamily)
// Port of medusalix/xone transport/dongle.c + the USB calls in mt76.c.

#include "usb/usb_transport.hpp"

#include <cerrno>
#include <condition_variable>
#include <cstdint>
#include <cstring>
#include <mutex>
#include <optional>
#include <thread>
#include <vector>

#include <CoreFoundation/CoreFoundation.h>
#include <IOKit/IOCFPlugIn.h>
#include <IOKit/IOKitLib.h>
#include <IOKit/usb/IOUSBLib.h>
#include <IOKit/usb/USB.h>

#include "common/log.hpp"

namespace xone::usb {

// Outstanding async reads per IN endpoint. The kernel-side driver uses 12
// URBs; user-space resubmits synchronously in the completion handler, so a
// small pool is enough to keep the pipes busy.
constexpr std::size_t num_in_reads = 4;

namespace {

// Numeric property matching on "idVendor"/"idProduct" is unreliable on some
// macOS versions (single-key matches fail while combined matches work), so
// we match the IOUSBDevice class and filter VID/PID in user space. The
// numeric values live in the "idVendor"/"idProduct" registry properties.
struct vid_pid {
    std::uint16_t vid = 0;
    std::uint16_t pid = 0;
};

auto read_vid_pid(io_service_t service) -> std::optional<vid_pid>
{
    CFMutableDictionaryRef props = nullptr;
    if (IORegistryEntryCreateCFProperties(service, &props, kCFAllocatorDefault, 0)
        != kIOReturnSuccess
        || !props)
        return std::nullopt;

    auto *vid_number = static_cast<CFNumberRef>(CFDictionaryGetValue(props, CFSTR("idVendor")));
    auto *pid_number = static_cast<CFNumberRef>(CFDictionaryGetValue(props, CFSTR("idProduct")));
    SInt32 v = 0, p = 0;
    std::optional<vid_pid> result;
    if (vid_number && pid_number
        && CFNumberGetValue(vid_number, kCFNumberSInt32Type, &v)
        && CFNumberGetValue(pid_number, kCFNumberSInt32Type, &p)) {
        result = { static_cast<std::uint16_t>(v), static_cast<std::uint16_t>(p) };
    }

    CFRelease(props);
    return result;
}

auto is_supported_pid(std::uint16_t pid) -> bool
{
    return pid == pid_old_dongle || pid == pid_new_dongle
        || pid == pid_builtin_asus_lenovo || pid == pid_surface_book_2;
}

auto is_dongle(vid_pid const &vp) -> bool
{
    return vp.vid == vid && is_supported_pid(vp.pid);
}

} // namespace

struct transport::read_slot {
    transport *owner = nullptr;
    std::uint8_t ep = 0;
    IOUSBInterfaceInterface190 **iface_ref = nullptr;
    UInt8 pipe_ref = 0;
    std::vector<std::uint8_t> buf;
};

struct transport::state {
    std::uint16_t pid = 0;

    // Device connection (port of the usb_device + usb_interface binding).
    // All interface references are double pointers: *ref is the interface
    // struct, and methods are called as (*ref)->Method(ref, ...).
    io_service_t service = 0;
    IOCFPlugInInterface **dev_iodev = nullptr;
    // Version 500 interface: adds ResetDevice (port of usb_reset_device).
    IOUSBDeviceInterface500 **dev_ref = nullptr;
    bool dev_opened = false;

    struct iface_conn {
        IOCFPlugInInterface **iodev = nullptr;
        IOUSBInterfaceInterface190 **iface_ref = nullptr;
    };
    std::vector<iface_conn> ifaces;
    bool re_enumerated = false;

    struct pipe_ref {
        IOUSBInterfaceInterface190 **iface_ref = nullptr;
        UInt8 ref = 0;
    };
    std::optional<pipe_ref> in_cmd_pipe;
    std::optional<pipe_ref> in_wlan_pipe;
    std::optional<pipe_ref> out_pipe;

    // One buffer per outstanding async read. Slots are never moved after the
    // initial ReadPipeAsync submissions (the refcon is their address).
    std::vector<read_slot> slots;

    // Async completion dispatch (one CFRunLoop on a dedicated thread).
    std::vector<CFRunLoopSourceRef> sources;
    IONotificationPortRef notify_port = 0;
    io_iterator_t termination_iter = 0;
    bool thread_started = false;

    std::thread thread;
    bool loop_ready = false;
    CFRunLoopRef runloop = nullptr;

    std::mutex lock;
    std::condition_variable cv;
    bool stopping = false;
    int in_flight = 0;
    bool disconnected_notified = false;

    frame_callback frames;
    disconnect_callback disconnected;
};

auto transport::probe(frame_callback frames, disconnect_callback disconnected) -> std::unique_ptr<transport>
{
    // make_unique cannot call the private constructor from outside the class.
    auto t = std::unique_ptr<transport>(new transport());
    if (!t->open(std::move(frames), std::move(disconnected)))
        return nullptr;
    return t;
}

void transport::stop_pump()
{
    {
        std::lock_guard<std::mutex> lock(m_state->lock);
        if (m_state->stopping)
            return;
        m_state->stopping = true;
    }

    for (auto &slot : m_state->slots)
        (*slot.iface_ref)->AbortPipe(slot.iface_ref, slot.pipe_ref);

    if (!m_state->thread_started)
        return;

    CFRunLoopRef runloop = nullptr;
    {
        std::unique_lock<std::mutex> lock(m_state->lock);
        m_state->cv.wait(lock, [this] { return m_state->loop_ready && m_state->in_flight == 0; });
        runloop = m_state->runloop;
    }

    // The reader thread is joined below; no callback can run after this.
    CFRunLoopStop(runloop);
    m_state->thread.join();
}

auto transport::re_enumerate() -> int
{
    // Do not use yet: a crashed SIE cannot finish the host-side port reset
    // handshake and drops off the bus. A crashed chip recovers on its own via
    // the MCU watchdog after ~90s of a quiet bus, so prefer waiting.
    // Stop the pump first so no completion callback touches the interface
    // references while the kernel tears them down.
    stop_pump();

    auto kr = (*m_state->dev_ref)->USBDeviceReEnumerate(m_state->dev_ref, kUSBAddExtraResetTimeMask);
    if (kr != kIOReturnSuccess) {
        xone::log_msg(log_level::error, "usb: re-enumerate failed (%d)", kr);
        return -EIO;
    }

    // The kernel terminated all of our clients; the USB references are dead.
    m_state->re_enumerated = true;
    xone::log_msg(log_level::info, "usb: re-enumerate ok");
    return 0;
}
transport::~transport()
{
    if (!m_state)
        return;

    stop_pump();

    // Release the termination watch and async event sources.
    if (m_state->termination_iter)
        IOObjectRelease(m_state->termination_iter);
    if (m_state->notify_port)
        IONotificationPortDestroy(m_state->notify_port);

    if (m_state->re_enumerated) {
        // The kernel already tore down the device and interfaces.
        if (m_state->service)
            IOObjectRelease(m_state->service);
        return;
    }

    for (auto *source : m_state->sources)
        CFRelease(source);

    // Close the interfaces and their endpoint pipes.
    for (auto &conn : m_state->ifaces) {
        (*conn.iface_ref)->USBInterfaceClose(conn.iface_ref);
        (*conn.iface_ref)->Release(conn.iface_ref);
        IODestroyPlugInInterface(conn.iodev);
    }

    // Close the device connection.
    if (m_state->dev_ref) {
        if (m_state->dev_opened)
            (*m_state->dev_ref)->USBDeviceClose(m_state->dev_ref);
        (*m_state->dev_ref)->Release(m_state->dev_ref);
    }
    if (m_state->dev_iodev)
        IODestroyPlugInInterface(m_state->dev_iodev);

    if (m_state->service)
        IOObjectRelease(m_state->service);
}

auto transport::open(frame_callback frames, disconnect_callback disconnected) -> bool
{
    m_state = std::make_unique<state>();
    m_state->frames = std::move(frames);
    m_state->disconnected = std::move(disconnected);

    // Find the dongle service (class match, filter by VID/PID).
    io_iterator_t iter = 0;
    kern_return_t kr = IOServiceGetMatchingServices(0, IOServiceMatching("IOUSBDevice"), &iter);
    if (kr != kIOReturnSuccess)
        return false;

    bool found = false;
    while (true) {
        io_service_t service = IOIteratorNext(iter);
        if (!service)
            break;
        auto vp = read_vid_pid(service);
        if (vp && is_dongle(*vp)) {
            m_state->service = service;
            m_state->pid = vp->pid;
            found = true;
        } else {
            IOObjectRelease(service);
        }
    }
    IOObjectRelease(iter);
    if (!found)
        return false;

    // Open the device connection.
    SInt32 score = 0;
    kr = IOCreatePlugInInterfaceForService(m_state->service, kIOUSBDeviceUserClientTypeID,
                                           kIOCFPlugInInterfaceID, &m_state->dev_iodev, &score);
    if (kr != kIOReturnSuccess) {
        xone::log_msg(log_level::error, "usb: create device interface failed (%d)", kr);
        return false;
    }

    void *slot = nullptr;
    HRESULT hr = (*m_state->dev_iodev)->QueryInterface(m_state->dev_iodev,
                                                      CFUUIDGetUUIDBytes(kIOUSBDeviceInterfaceID500), &slot);
    if (hr != S_OK || !slot) {
        xone::log_msg(log_level::error, "usb: query device interface failed");
        return false;
    }
    m_state->dev_ref = static_cast<IOUSBDeviceInterface500 **>(slot);

    kr = (*m_state->dev_ref)->USBDeviceOpen(m_state->dev_ref);
    if (kr != kIOReturnSuccess) {
        xone::log_msg(log_level::error, "usb: open device failed (%d)", kr);
        return false;
    }
    m_state->dev_opened = true;

    // Reset the chip so it starts from the boot ROM (port of the
    // usb_reset_device call in xone_dongle_probe).
    kr = (*m_state->dev_ref)->ResetDevice(m_state->dev_ref);
    if (kr != kIOReturnSuccess) {
        xone::log_msg(log_level::error, "usb: reset device failed (%d)", kr);
        return false;
    }

    // Open every interface and collect the endpoint pipes we need.
    io_iterator_t children = 0;
    kr = IORegistryEntryGetChildIterator(m_state->service, kIOServicePlane, &children);
    if (kr != kIOReturnSuccess)
        return false;

    while (true) {
        io_service_t child = IOIteratorNext(children);
        if (!child)
            break;
        open_interface(child);
        IOObjectRelease(child);
    }
    IOObjectRelease(children);

    if (!m_state->in_cmd_pipe || !m_state->in_wlan_pipe || !m_state->out_pipe) {
        xone::log_msg(log_level::error, "usb: missing endpoint (cmd in=%d, wlan in=%d, out=%d)",
                      m_state->in_cmd_pipe.has_value(), m_state->in_wlan_pipe.has_value(),
                      m_state->out_pipe.has_value());
        return false;
    }

    // Async completion dispatch for the opened interfaces.
    for (auto &conn : m_state->ifaces) {
        CFRunLoopSourceRef source = nullptr;
        if ((*conn.iface_ref)->CreateInterfaceAsyncEventSource(conn.iface_ref, &source)
            == kIOReturnSuccess
            && source)
            m_state->sources.push_back(source);
    }

    // Watch for the dongle going away (unplug or chip reconnect). The
    // matching dictionary is consumed by this call.
    m_state->notify_port = IONotificationPortCreate(0);
    kr = IOServiceAddMatchingNotification(m_state->notify_port, kIOTerminatedNotification,
                                          IOServiceMatching("IOUSBDevice"),
                                          on_dongle_terminated, this,
                                          &m_state->termination_iter);
    if (kr != kIOReturnSuccess) {
        xone::log_msg(log_level::error, "usb: add termination notification failed (%d)", kr);
        return false;
    }

    m_state->thread = std::thread([this] { worker_loop(); });
    m_state->thread_started = true;

    // Submit the initial async reads (EP 0x05 IN and EP 0x04 IN). Reserve so
    // the slot addresses stay valid for the ReadPipeAsync refcons.
    m_state->slots.reserve(num_in_reads * 2);
    auto submit_all = [this](std::optional<state::pipe_ref> const &pipe,
                             std::uint8_t ep, std::size_t buf_len) -> bool {
        for (std::size_t i = 0; i < num_in_reads; i++) {
            m_state->slots.push_back(
                { this, ep, pipe->iface_ref, pipe->ref, std::vector<std::uint8_t>(buf_len) });
            if (!submit_read(&m_state->slots.back()))
                return false;
        }
        return true;
    };

    if (!submit_all(m_state->in_cmd_pipe, ep_in_cmd, len_cmd_pkt))
        return false;
    if (!submit_all(m_state->in_wlan_pipe, ep_in_wlan, len_wlan_pkt))
        return false;

    xone::log_msg(log_level::info, "usb: dongle connected (pid=0x%04x)", m_state->pid);
    return true;
}

auto transport::open_interface(io_service_t child) -> bool
{
    state::iface_conn conn{};
    SInt32 score = 0;
    if (IOCreatePlugInInterfaceForService(child, kIOUSBInterfaceUserClientTypeID,
                                          kIOCFPlugInInterfaceID, &conn.iodev, &score)
        != kIOReturnSuccess)
        return false;

    void *slot = nullptr;
    if ((*conn.iodev)->QueryInterface(conn.iodev, CFUUIDGetUUIDBytes(kIOUSBInterfaceInterfaceID190),
                                      &slot) != S_OK
        || !slot) {
        IODestroyPlugInInterface(conn.iodev);
        return false;
    }
    conn.iface_ref = static_cast<IOUSBInterfaceInterface190 **>(slot);

    if ((*conn.iface_ref)->USBInterfaceOpen(conn.iface_ref) != kIOReturnSuccess) {
        (*conn.iface_ref)->Release(conn.iface_ref);
        IODestroyPlugInInterface(conn.iodev);
        return false;
    }

    m_state->ifaces.push_back(conn);

    // Scan the interface's pipes for the endpoints we need.
    UInt8 num_pipes = 0;
    if ((*conn.iface_ref)->GetNumEndpoints(conn.iface_ref, &num_pipes) != kIOReturnSuccess)
        return true;

    for (UInt8 i = 1; i <= num_pipes; i++) {
        UInt8 direction = 0, number = 0, type = 0;
        UInt16 max_packet_size = 0;
        UInt8 interval = 0;
        if ((*conn.iface_ref)->GetPipeProperties(conn.iface_ref, i, &direction, &number, &type,
                                                 &max_packet_size, &interval) != kIOReturnSuccess)
            continue;

        std::uint8_t ep = static_cast<std::uint8_t>(number);
        if (direction == kUSBIn && ep == ep_in_cmd && !m_state->in_cmd_pipe.has_value())
            m_state->in_cmd_pipe = { conn.iface_ref, i };
        else if (direction == kUSBIn && ep == ep_in_wlan && !m_state->in_wlan_pipe.has_value())
            m_state->in_wlan_pipe = { conn.iface_ref, i };
        else if (direction == kUSBOut && ep == ep_out && !m_state->out_pipe.has_value())
            m_state->out_pipe = { conn.iface_ref, i };
    }

    return true;
}

auto transport::pid() const -> std::uint16_t
{
    return m_state->pid;
}

auto transport::send_vendor_request(vendor_request req, bool is_read, std::uint16_t w_value,
                                    std::uint16_t w_index, void *data, std::size_t len) -> int
{
    IOUSBDevRequest request;
    std::memset(&request, 0, sizeof(request));
    request.bmRequestType = USBmakebmRequestType(is_read ? static_cast<int>(kUSBIn) : static_cast<int>(kUSBOut),
                                                  static_cast<int>(kUSBVendor), static_cast<int>(kUSBDevice));
    request.bRequest = static_cast<UInt8>(req);
    request.wValue = w_value;
    request.wIndex = w_index;
    request.wLength = static_cast<UInt16>(len);
    request.pData = data;

    IOReturn ret = (*m_state->dev_ref)->DeviceRequest(m_state->dev_ref, &request);
    if (ret != kIOReturnSuccess || request.wLenDone != len) {
        xone::log_msg(log_level::error, "usb: vendor request 0x%02x failed (%d)",
                      static_cast<int>(req), ret);
        return -EIO;
    }

    return static_cast<int>(len);
}

auto transport::bulk_write(void const *data, std::size_t len) -> int
{
    auto &pipe = m_state->out_pipe.value();
    IOReturn ret = (*pipe.iface_ref)->WritePipe(pipe.iface_ref, pipe.ref,
                                                 const_cast<void *>(data), static_cast<UInt32>(len));
    if (ret != kIOReturnSuccess) {
        xone::log_msg(log_level::error, "usb: bulk write failed (%d)", ret);
        return -EIO;
    }

    return static_cast<int>(len);
}

auto transport::submit_read(read_slot *slot) -> bool
{
    {
        std::lock_guard<std::mutex> lock(m_state->lock);
        if (m_state->stopping)
            return false;
        m_state->in_flight++;
    }

    IOReturn ret = (*slot->iface_ref)->ReadPipeAsync(slot->iface_ref, slot->pipe_ref,
                                                     slot->buf.data(), static_cast<UInt32>(slot->buf.size()),
                                                     on_read_completion, slot);
    if (ret != kIOReturnSuccess) {
        std::lock_guard<std::mutex> lock(m_state->lock);
        m_state->in_flight--;
        m_state->cv.notify_all();
        return false;
    }

    return true;
}

void transport::worker_loop()
{
    CFRunLoopRef runloop = CFRunLoopGetCurrent();
    for (auto *source : m_state->sources)
        CFRunLoopAddSource(runloop, source, kCFRunLoopDefaultMode);
    CFRunLoopAddSource(runloop, IONotificationPortGetRunLoopSource(m_state->notify_port),
                       kCFRunLoopDefaultMode);

    {
        std::lock_guard<std::mutex> lock(m_state->lock);
        m_state->runloop = runloop;
        m_state->loop_ready = true;
        m_state->cv.notify_all();
    }

    CFRunLoopRun();
}

void transport::handle_read(read_slot *slot, IOReturn result, std::size_t len)
{
    if (result == kIOReturnSuccess && len > 0 && m_state->frames)
        m_state->frames(slot->ep, slot->buf.data(), len);

    std::lock_guard<std::mutex> lock(m_state->lock);
    if (m_state->stopping) {
        m_state->in_flight--;
        m_state->cv.notify_all();
        return;
    }

    IOReturn ret = (*slot->iface_ref)->ReadPipeAsync(slot->iface_ref, slot->pipe_ref,
                                                     slot->buf.data(), static_cast<UInt32>(slot->buf.size()),
                                                     on_read_completion, slot);
    if (ret != kIOReturnSuccess) {
        m_state->in_flight--;
        m_state->cv.notify_all();
    }
}

void transport::handle_disconnected()
{
    disconnect_callback callback;
    {
        std::lock_guard<std::mutex> lock(m_state->lock);
        if (m_state->stopping || m_state->disconnected_notified)
            return;
        m_state->disconnected_notified = true;
        callback = m_state->disconnected;
    }

    xone::log_msg(log_level::info, "usb: dongle disconnected");
    if (callback)
        callback();
}

void transport::on_read_completion(void *refcon, IOReturn result, void *arg0)
{
    auto *slot = static_cast<read_slot *>(refcon);
    slot->owner->handle_read(slot, result,
                             static_cast<std::size_t>(reinterpret_cast<std::uintptr_t>(arg0)));
}

void transport::on_dongle_terminated(void *refcon, io_iterator_t iter)
{
    auto *t = static_cast<transport *>(refcon);

    bool ours = false;
    while (true) {
        io_service_t service = IOIteratorNext(iter);
        if (!service)
            break;
        auto vp = read_vid_pid(service);
        if (vp && vp->vid == vid && vp->pid == t->m_state->pid)
            ours = true;
        IOObjectRelease(service);
    }

    if (ours)
        t->handle_disconnected();
}

auto dongle_present() -> bool
{
    io_iterator_t iter = 0;
    kern_return_t kr = IOServiceGetMatchingServices(0, IOServiceMatching("IOUSBDevice"), &iter);
    if (kr != kIOReturnSuccess)
        return false;

    bool present = false;
    while (true) {
        io_service_t service = IOIteratorNext(iter);
        if (!service)
            break;
        auto vp = read_vid_pid(service);
        if (vp && is_dongle(*vp))
            present = true;
        IOObjectRelease(service);
    }
    IOObjectRelease(iter);
    return present;
}

} // namespace xone::usb