#include <stdbool.h>
#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <err.h>
#include "hidpp10.h"
#include "hidpp20.h"

size_t read_fw(const char *filename, uint8_t **fw_buf) {
    int fw_fd;
    ssize_t r;
    size_t read_bytes = 0;
    size_t bufsize = 0;
    const size_t blocksize = 1024;

    fw_fd = open(filename, O_RDONLY);
    if (fw_fd < 0) {
        perror(filename);
        return 0;
    }

    *fw_buf = NULL;
    do {
        if (bufsize - read_bytes < blocksize) {
            bufsize += blocksize;
            *fw_buf = realloc(*fw_buf, bufsize);
            if (fw_buf == NULL) {
                perror("realloc");
                abort();
            }
        }

        r = read(fw_fd, *fw_buf + read_bytes, bufsize - read_bytes);
        if (r > 0)
            read_bytes += (size_t) r;
    } while (r > 0);

    if (r < 0)
        perror("read");
    if (read_bytes == 0)
        fprintf(stderr, "Warning: no data is read\n");

    close(fw_fd);

    if (r < 0 || read_bytes == 0) {
        free(*fw_buf);
        *fw_buf = NULL;
        return 0;
    }

    return read_bytes;
}

#if 1
#define DPRINTF(...) printf(__VA_ARGS__)
#endif

/* HID++-specific routines */
#define WPID_T650 0x4101 /* Wireless PID of T650 touchpad */
#define FEATURE_ID_DFU 0x00C0

struct match_notif_data {
    uint8_t report_id;
    uint8_t device_index;
    uint8_t notif_id;
    bool matched;
    HidppMessage match;
};
static bool match_notif(HidppMessage *msg, void *userdata) {
    struct match_notif_data *d = userdata;
    DPRINTF("Received report: %02x %02x %02x  %02x %02x %02x %02x\n",
        msg->report_id, msg->device_index, msg->sub_id,
        msg->address, msg->params[0], msg->params[1], msg->params[2]);
    if (msg->report_id == d->report_id &&
        msg->device_index == d->device_index &&
        msg->sub_id == d->notif_id) {
        d->matched = true;
        d->match = *msg;
        return true;
    }
    return false;
}

#if 0
#define RETURN_IF_FAIL(cond, ...) \
    if (!(cond)) { warnx(__VA_ARGS__); return false; }
#else
/* debug: additionally print the message */
#define RETURN_IF_FAIL(cond, ...) \
    if (!(cond)) { warnx(__VA_ARGS__); return false; } \
    else { DPRINTF("PASSED: " __VA_ARGS__); DPRINTF("\n"); }
#define PRWARN_IF_FAIL(cond, ...) \
    if (!(cond)) { warnx(__VA_ARGS__); } \
    else { DPRINTF("PASSED: " __VA_ARGS__); DPRINTF("\n"); }
#endif

bool fw_update(int fd, uint8_t device_index, uint8_t *fw, size_t fw_len) {
    int r;
    int i;
    bool b;
    uint16_t version;
    FeatureInfo feat_dfu;
    uint8_t *params;
    uint16_t wpid;

    // [ix feat] find HID++ version, confirm >= 0x0200
    version = hidpp20_get_version(fd, device_index);
    PRWARN_IF_FAIL(version >= 0x0200, "HID++ 2.0 wanted, got %#x", version);

    // [ix feat] discover DFU feature
    r = hidpp20_get_feature_by_id(fd, device_index, FEATURE_ID_DFU, &feat_dfu);
    PRWARN_IF_FAIL(r == 0, "DFU feature not available");
    if (r != 0) { /* HACK: hard-coded for specific T650  */
        feat_dfu.feature_index = 13;
        feat_dfu.feature_id = FEATURE_ID_DFU;
        feat_dfu.feature_type = 0;
    }

    r = hidpp10_enable_wireless_notifications(fd, true);
    RETURN_IF_FAIL(r == 0, "Failed to enable wireless notifications!");

    // print: ask user to restart device
    puts("Please restart your device");

    // [ix notif] wait for device connect
    HidppMessage conn_notif;
    struct match_notif_data match_notif_data = {
        .report_id      = HIDPP_SHORT,
        .device_index   = device_index,
        .notif_id       = NOTIF_DEV_CONNECT
    };
    b = hidpp_read_msg(fd, 10000, &conn_notif, match_notif, &match_notif_data);
    params = conn_notif.params;
    PRWARN_IF_FAIL(b, "Timeout waiting for device connect");
    PRWARN_IF_FAIL(params[0] == 4, "Connected device is not Unifying");
    PRWARN_IF_FAIL(params[1] & (1 << 6), "Device is not connected");
    wpid = (params[2] << 8) | params[1];
    PRWARN_IF_FAIL (wpid == WPID_T650, "Not a T650 (%04x) but WPID:%04x",
        WPID_T650, wpid);

    // [ix feat req] send "DFU" to DFU feature
    HidppMessage req_dfu = {
        .report_id      = HIDPP_SHORT,
        .device_index   = device_index,
        .feature_index  = feat_dfu.feature_index,
        .func           = HIDPP20_FUNC(0),
    };
    r = hidpp10_request(fd, &req_dfu, NULL, NULL);
    PRWARN_IF_FAIL(r == 0, "DFU feature not found!");
    // [ix feat resp] confirm DFU enablement
    PRWARN_IF_FAIL(req_dfu.params[0] == 0, "DFU not enabled?");

    // [ff SHORT_REG req] write "LT" ix to f0
    const uint8_t lt_params[3] = { 'L', 'T', device_index };
    HidppMessage req_set_lt = {
        .report_id      = HIDPP_SHORT,
        .device_index   = 0xFF, /* receiver */
        .sub_id         = SUB_SET_REGISTER,
        .address        = 0xF0,
    };
    memcpy(req_set_lt.params, lt_params, 3);
    struct match_notif_data req_lt_notif = {
        .report_id      = DJ_SHORT,
        .device_index   = device_index,
        .notif_id       = 0x42, /* Connection Status */
    };
    r = hidpp10_request(fd, &req_set_lt, match_notif, &req_lt_notif);
#if 0
    // [DJ ix notif resp] Link-loss notif (0x42)
    /* maybe this doesn't work because DJ notifs are not enabled? */
    RETURN_IF_FAIL(req_lt_notif.matched, "Did not receive link-loss notif!");
    RETURN_IF_FAIL(req_lt_notif.match.params[0] & 1, "Expected link-loss!");
#endif
    // [ff SHORT_REG resp] write must be succesful
    RETURN_IF_FAIL(r == 0, "SET_REG(0xF0, \"LT\\%i\") failed with %i\n",
            device_index, r);

    // [ff SHORT_REG] read f0 and confirm it says LT ix
    HidppMessage req_get_lt = {
        .report_id      = HIDPP_SHORT,
        .device_index   = 0xFF, /* receiver */
        .sub_id         = SUB_GET_REGISTER,
        .address        = 0xF0,
    };
    r = hidpp10_request(fd, &req_get_lt, NULL, NULL);
    RETURN_IF_FAIL(r == 0, "GET_REG(0xF0) failed with %i\n", r);
    RETURN_IF_FAIL(memcmp(req_get_lt.params, lt_params, 3) == 0,
        "Expected \"LT\\%i\" from 0xF0 register!", device_index);
    /* Now "DUD communication" is acknowledged */

#define DFU_INIT(code, ...) \
    ((HidppMessage) { \
        .report_id      = HIDPP_LONG, \
        .device_index   = device_index, \
        .sub_id         = SUB_SET_LONG_REGISTER, \
        .address        = 0xE2, \
        .params_l       = { code, __VA_ARGS__ } \
    })
#define DFU_WRITE(_msg) \
    do { \
        r = hidpp10_request(fd, &_msg, NULL, NULL); \
        RETURN_IF_FAIL(r == 0, "Write E2 failed"); \
    } while (0)

    // TODO: [ix 82 E2] send magic packet (+verify)
    HidppMessage dfu_req;
    // 0X11 0X2 0X82 0Xe2 0X2 0X9 0X5 0X41 0X1 0X0 0X38 0X0 0X0 0X0 0X0 0X0
    // 11 02 82 e2  HID++ long, device index, SubID (Set_Reg_Long), Addr
    // 02       DFU upgrade signature packet command?
    // 09 05    ???
    // 41 01    Wireless PID (MSB LSB)
    // 00 38    New firmware version (build number)
    // 00 00 00 00 00 (padding)
    wpid = WPID_T650;
    for (i = 3; i >= 0; i--) {
        RETURN_IF_FAIL(i > 1, "Too many DFU magic attempts");
        dfu_req = DFU_INIT(0x02, 0x09, 0x05,
                           wpid >> 8, (uint8_t) wpid,
                           0x00, 0x38);
        DFU_WRITE(dfu_req);
        /* verify response */
        int code = dfu_req.params[0];
        if (code == 0x12) {         /* Retry? */
        } else if (code == 0x01) {  /* toggled? */
            /* assume that 0x12 wants us to retry, and 1 means success. */
            break;
        } else {
            RETURN_IF_FAIL(code == 0x12 || code == 0x01,
                    "Unexpected DFU resp %#04x", code);
        }
    }

    // TODO: [ix 82 E2] send firmware (+check toggle)
    (void)fw, (void)fw_len;

    // TODO: [ix 82 E2] send verify (+check toggle)
    dfu_req = DFU_INIT(0x04);
    DFU_WRITE(dfu_req);
    RETURN_IF_FAIL(dfu_req.params[0] == 0x01, "DFU verification failure?");

    // TODO: [ff SHORT_REG req] send "XIT
    HidppMessage xit_req = {
        .report_id      = HIDPP_SHORT,
        .device_index   = 0xFF,
        .sub_id         = SUB_SET_REGISTER,
        .address        = 0xF0,
        .params         = { 'X', 'I', 'T' }
    };
    struct match_notif_data recon_data = {
        .report_id      = HIDPP_SHORT,
        .device_index   = device_index,
        .notif_id       = NOTIF_DEV_CONNECT
    };
    hidpp10_request(fd, &xit_req, match_notif, &recon_data);
    RETURN_IF_FAIL(r == 0, "XIT failed");

    // TODO: [ix notif] expect re-connection (0x41, proto 5)
    if (!recon_data.matched) {
        b = hidpp_read_msg(fd, 10000, &conn_notif, match_notif, &recon_data);
        RETURN_IF_FAIL(b, "Timeout waiting for re-connect");
    }
    /* not addr, but notif has different formatting */
    uint8_t proto = recon_data.match.address;
    RETURN_IF_FAIL(proto == 5, "Expected proto 5, got %#04x", proto);
    params = recon_data.match.params_l;
    RETURN_IF_FAIL(params[0] & (1 << 6), "Device is not connected");
    wpid = (params[2] << 8) | params[1];
    RETURN_IF_FAIL (wpid == WPID_T650, "Not a T650 (%#04x) but WPID:%04x",
        WPID_T650, wpid);

    // TODO: (optional): ask user to restart device, confirm notif
    puts("Please restart your device");

    // print: success
    puts("Firmware update success!");

    return true;
}

int main(int argc, char **argv) {
    uint8_t *fw = NULL;
    size_t fw_len;
    int hid_fd = -1;
    int r = 0;
    uint8_t device_index;

    if (argc <= 3) {
        fprintf(stderr, "Usage: %s /dev/hidrawX device_index flash-file\n", argv[0]);
        return 1;
    }

    device_index = (uint8_t) (argv[2][0] - '0');
    if (!(device_index >= 1 && device_index <= 6)) {
        fprintf(stderr, "device_index must be between 1 and 6 (inclusive)\n");
        return 1;
    }

    if ((fw_len = read_fw(argv[3], &fw)) == 0) {
        return 1;
    }

    printf("Firmware size: %zu\n", fw_len);
    /* ensure that firmware is a multiple of 15 bytes */
    if (fw_len % 15) {
        size_t fw_len_padded = (fw_len / 15 + 1) * 15;
        fw = realloc(fw, fw_len_padded);
        if (fw == NULL) {
            perror("realloc");
            abort();
        }
        memset(fw + fw_len, 0xff, fw_len_padded - fw_len);
        printf("  (padded with %zi bytes)\n", fw_len_padded - fw_len);
        fw_len = fw_len_padded;
    }

    hid_fd = open(argv[1], O_RDWR);
    if (hid_fd >= 0) {
        r = fw_update(hid_fd, device_index, fw, fw_len) ? 0 : 1;
    } else {
        perror(argv[1]);
        r = 1;
    }

    if (hid_fd >= 0)
        close(hid_fd);
    free(fw);
    return r;
}

/* vim: set sw=4 et ts=4: */