#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;
    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); \
        PRWARN_IF_FAIL(r == 0, "Write E2 failed, r=%#04x", r); \
    } while (0)

    HidppMessage dfu_req;
    /* Format (req): E2 [action] [firmware blob]
     * Possible action for requests (from sw to dvc; it seems to be dependent on
     * the responses received from the device):
     *      02  indicates the begin of firmware (always sent by sw if the DFU
     *          Upgrade Signature bit (rsp.bit1) is set by dvc). ("even toggle
     *          bit" as it is the first fw block at index 0)
     *   00/01  for even/odd fw block indices.
     *      04  DFU failed, execute recovery procedure. (params are all 00)
     *      08  DFU end (verify?). (params are all 00)
     * If rsp.bit2 is set, then respond with the same result code and firmware
     * blob in place of the action. Probably for this reason, the device sends
     * result 12 instead of just 02 which would have an ambiguous meaning.
     *
     * Format (rsp): E2 [result]
     * Sw expects the result to toggle, so if action is 01, then result is
     * 00 and if action is 00, then result is 01. Bit layout for rsp:
     * 7-4  f0  ignored (but hw must copy this)
     *   3  08  set to signal a critical error (sw must abort with action 04)
     *   2  04  Unidentified packet failure (sw must retry same fw data)
     *   1  02  DFU Upgrade signature bit (sw must start with begin of fw)
     *   0  01  toggle bit (dvc flips this bit for every response. If it
     *          doesn't, then Sw should warn and expect the next bit to be the
     *          complement of this one)
     */
    unsigned retries = 0;
    // [ix 82 E2] send magic packet (+verify)
    uint8_t action = 2;
    uint8_t dvc_toggle = -1;
    size_t fw_pos = 0;
    printf("Starting DFU (%zi blocks)\n", fw_len / 15);
    while (fw_pos < fw_len) {
        // [ix 82 E2] send firmware (+check toggle)
        dfu_req = DFU_INIT(action);
        memcpy(&dfu_req.params_l[1], fw + fw_pos, 15);
        r = hidpp10_request(fd, &dfu_req, NULL, NULL);
        if (r != 0) {
            /* the official updater does 4 attempts (3 retries), this is
             * simpler. */
            warnx("Device does not respond! Aborting");
            break;
        }

        /* verify response */
        int code = dfu_req.params[0];
        PRWARN_IF_FAIL(dvc_toggle != (uint8_t) -1 && (code & 1) == dvc_toggle,
            "Device did not toggle properly!");
        dvc_toggle = (code & 1) ^ 1; /* toggle flag: inverse of previous resp */

        if (code & 8) { /* Critical Error, abort! */
            warnx("The Device does not accept this firmware code (Wrong Signature)");
            break;
        } else if (code & 2) { /* DFU Upgrade signature, reset! */
            if (retries++ < 12) {
                warnx("DFU error at pos %#zx, retry %i\n", fw_pos, retries);
                dvc_toggle = -1;
                fw_pos = 0;
                action = 2;
            } else {
                warnx("DFU error, not retrying anymore!");
                break;
            }
        } else if (code & 4) { /* Unidentified packet failure, resend packet */
            action = code; /* send this "tag"/result as "tag"/action */
            /* TODO: maybe check for infinite loop? The official fw updater
             * doesn't do this either though, stopped counting after 61
             * retries... */
        } else { /* fw packet ack. */
            action = (fw_pos / 15 + 1) & 1; /* toggle bit for odd blocks */
            fw_pos += 15; /* go to next fw block */
        }
    }

    if (fw_pos == fw_len) { /* DFU success */
        // [ix 82 E2] send verify (+check toggle)
        dfu_req = DFU_INIT(8); /* verify, write, or whatever */
        DFU_WRITE(dfu_req);
        PRWARN_IF_FAIL((dfu_req.params[0] & ~1) == 0x04,
                "DFU verification failure?");
    } else { /* DFU failed */
        warnx("DFU failed, executing recovery procedure");
        dfu_req = DFU_INIT(4);
        DFU_WRITE(dfu_req);
    }

    // [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, "Failed to deactivate DFU mode");

    // [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: */