// SPDX-License-Identifier: GPL-2.0-only
/* Copyright (c) 2024 Red Hat, Inc
 */

#include "vmlinux.h"
#include "hid_bpf.h"
#include "hid_bpf_helpers.h"
#include "hid_report_helpers.h"
#include <bpf/bpf_tracing.h>

#define HID_BPF_ASYNC_MAX_CTX 1
#include "hid_bpf_async.h"

#define VID_UGEE		0x28BD
/* same PID whether connected directly or through the provided dongle: */
#define PID_ACK05_REMOTE	0x0202


HID_BPF_CONFIG(
	HID_DEVICE(BUS_USB, HID_GROUP_GENERIC, VID_UGEE, PID_ACK05_REMOTE),
);

/*
 * By default, the pad reports the buttons through a set of key sequences.
 *
 * The pad reports a classic keyboard report descriptor:
 * # HANVON UGEE Shortcut Remote
 * Report descriptor length: 102 bytes
 *  0x05, 0x01,                    // Usage Page (Generic Desktop)              0
 *  0x09, 0x02,                    // Usage (Mouse)                             2
 *  0xa1, 0x01,                    // Collection (Application)                  4
 *  0x85, 0x09,                    //   Report ID (9)                           6
 *  0x09, 0x01,                    //   Usage (Pointer)                         8
 *  0xa1, 0x00,                    //   Collection (Physical)                   10
 *  0x05, 0x09,                    //     Usage Page (Button)                   12
 *  0x19, 0x01,                    //     UsageMinimum (1)                      14
 *  0x29, 0x03,                    //     UsageMaximum (3)                      16
 *  0x15, 0x00,                    //     Logical Minimum (0)                   18
 *  0x25, 0x01,                    //     Logical Maximum (1)                   20
 *  0x95, 0x03,                    //     Report Count (3)                      22
 *  0x75, 0x01,                    //     Report Size (1)                       24
 *  0x81, 0x02,                    //     Input (Data,Var,Abs)                  26
 *  0x95, 0x05,                    //     Report Count (5)                      28
 *  0x81, 0x01,                    //     Input (Cnst,Arr,Abs)                  30
 *  0x05, 0x01,                    //     Usage Page (Generic Desktop)          32
 *  0x09, 0x30,                    //     Usage (X)                             34
 *  0x09, 0x31,                    //     Usage (Y)                             36
 *  0x26, 0xff, 0x7f,              //     Logical Maximum (32767)               38
 *  0x95, 0x02,                    //     Report Count (2)                      41
 *  0x75, 0x10,                    //     Report Size (16)                      43
 *  0x81, 0x02,                    //     Input (Data,Var,Abs)                  45
 *  0x05, 0x0d,                    //     Usage Page (Digitizers)               47
 *  0x09, 0x30,                    //     Usage (Tip Pressure)                  49
 *  0x26, 0xff, 0x07,              //     Logical Maximum (2047)                51
 *  0x95, 0x01,                    //     Report Count (1)                      54
 *  0x75, 0x10,                    //     Report Size (16)                      56
 *  0x81, 0x02,                    //     Input (Data,Var,Abs)                  58
 *  0xc0,                          //   End Collection                          60
 *  0xc0,                          // End Collection                            61
 *  0x05, 0x01,                    // Usage Page (Generic Desktop)              62
 *  0x09, 0x06,                    // Usage (Keyboard)                          64
 *  0xa1, 0x01,                    // Collection (Application)                  66
 *  0x85, 0x06,                    //   Report ID (6)                           68
 *  0x05, 0x07,                    //   Usage Page (Keyboard/Keypad)            70
 *  0x19, 0xe0,                    //   UsageMinimum (224)                      72
 *  0x29, 0xe7,                    //   UsageMaximum (231)                      74
 *  0x15, 0x00,                    //   Logical Minimum (0)                     76
 *  0x25, 0x01,                    //   Logical Maximum (1)                     78
 *  0x75, 0x01,                    //   Report Size (1)                         80
 *  0x95, 0x08,                    //   Report Count (8)                        82
 *  0x81, 0x02,                    //   Input (Data,Var,Abs)                    84
 *  0x05, 0x07,                    //   Usage Page (Keyboard/Keypad)            86
 *  0x19, 0x00,                    //   UsageMinimum (0)                        88
 *  0x29, 0xff,                    //   UsageMaximum (255)                      90
 *  0x26, 0xff, 0x00,              //   Logical Maximum (255)                   92
 *  0x75, 0x08,                    //   Report Size (8)                         95
 *  0x95, 0x06,                    //   Report Count (6)                        97
 *  0x81, 0x00,                    //   Input (Data,Arr,Abs)                    99
 *  0xc0,                          // End Collection                            101
 *
 * Each button gets assigned the following events:
 *
 *   Buttons released: 06 00 00 00 00 00 00 00
 *   Button 1:         06 01 12 00 00 00 00 00 -> LControl + o
 *   Button 2:         06 01 11 00 00 00 00 00 -> LControl + n
 *   Button 3:         06 00 3e 00 00 00 00 00 -> F5
 *   Button 4:         06 02 00 00 00 00 00 00 -> LShift
 *   Button 5:         06 01 00 00 00 00 00 00 -> LControl
 *   Button 6:         06 04 00 00 00 00 00 00 -> LAlt
 *   Button 7:         06 01 16 00 00 00 00 00 -> LControl + s
 *   Button 8:         06 01 1d 00 00 00 00 00 -> LControl + z
 *   Button 9:         06 00 2c 00 00 00 00 00 -> Space
 *   Button 10:        06 03 1d 00 00 00 00 00 -> LControl + LShift + z
 *   Wheel:            06 01 57 00 00 00 00 00 -> clockwise rotation (LControl + Keypad Plus)
 *   Wheel:            06 01 56 00 00 00 00 00 -> counter-clockwise rotation
 *						  (LControl + Keypad Minus)
 *
 * However, multiple buttons can be pressed at the same time, and when this happens,
 * each button gets assigned a new slot in the Input (Data,Arr,Abs):
 *
 *   Button 1 + 3:     06 01 12 3e 00 00 00 00 -> LControl + o + F5
 *
 * When a modifier is pressed (Button 4, 5, or 6), the assigned key is set to 00:
 *
 *   Button 5 + 7:     06 01 00 16 00 00 00 00 -> LControl + s
 *
 * This is mostly fine, but with Button 8 and Button 10 sharing the same
 * key value ("z"), there are cases where we can not know which is which.
 *
 */

#define PAD_WIRED_DESCRIPTOR_LENGTH 102
#define PAD_DONGLE_DESCRIPTOR_LENGTH 177
#define STYLUS_DESCRIPTOR_LENGTH 109
#define VENDOR_DESCRIPTOR_LENGTH 36
#define PAD_REPORT_ID 6
#define RAW_PAD_REPORT_ID 0xf0
#define RAW_BATTERY_REPORT_ID 0xf2
#define VENDOR_REPORT_ID 2
#define PAD_REPORT_LENGTH 8
#define VENDOR_REPORT_LENGTH 12

__u16 last_button_state;

static const __u8 disabled_rdesc[] = {
	// Make sure we match our original report length
	FixedSizeVendorReport(VENDOR_REPORT_LENGTH)
};

static const __u8 fixed_rdesc_vendor[] = {
	UsagePage_GenericDesktop
	Usage_GD_Keypad
	CollectionApplication(
		// -- Byte 0 in report
		ReportId(RAW_PAD_REPORT_ID)
		// Byte 1 in report - same than report ID
		ReportCount(1)
		ReportSize(8)
		Input(Const) // padding (internal report ID)
		LogicalMaximum_i8(0)
		LogicalMaximum_i8(1)
		UsagePage_Digitizers
		Usage_Dig_TabletFunctionKeys
		CollectionPhysical(
			// Byte 2-3 is the button state
			UsagePage_Button
			UsageMinimum_i8(0x01)
			UsageMaximum_i8(0x0a)
			LogicalMinimum_i8(0x0)
			LogicalMaximum_i8(0x1)
			ReportCount(10)
			ReportSize(1)
			Input(Var|Abs)
			Usage_i8(0x31) // will be mapped as BTN_A / BTN_SOUTH
			ReportCount(1)
			Input(Var|Abs)
			ReportCount(5) // padding
			Input(Const)
			// Byte 4 in report - just exists so we get to be a tablet pad
			UsagePage_Digitizers
			Usage_Dig_BarrelSwitch // BTN_STYLUS
			ReportCount(1)
			ReportSize(1)
			Input(Var|Abs)
			ReportCount(7) // padding
			Input(Const)
			// Bytes 5/6 in report - just exists so we get to be a tablet pad
			UsagePage_GenericDesktop
			Usage_GD_X
			Usage_GD_Y
			ReportCount(2)
			ReportSize(8)
			Input(Var|Abs)
			// Byte 7 in report is the dial
			Usage_GD_Wheel
			LogicalMinimum_i8(-1)
			LogicalMaximum_i8(1)
			ReportCount(1)
			ReportSize(8)
			Input(Var|Rel)
		)
		// -- Byte 0 in report
		ReportId(RAW_BATTERY_REPORT_ID)
		// Byte 1 in report - same than report ID
		ReportCount(1)
		ReportSize(8)
		Input(Const) // padding (internal report ID)
		// Byte 2 in report - always 0x01
		Input(Const) // padding (internal report ID)
		UsagePage_Digitizers
		/*
		 * We represent the device as a stylus to force the kernel to not
		 * directly query its battery state. Instead the kernel will rely
		 * only on the provided events.
		 */
		Usage_Dig_Stylus
		CollectionPhysical(
			// Byte 3 in report - battery value
			UsagePage_BatterySystem
			Usage_BS_AbsoluteStateOfCharge
			LogicalMinimum_i8(0)
			LogicalMaximum_i8(100)
			ReportCount(1)
			ReportSize(8)
			Input(Var|Abs)
			// Byte 4 in report - charging state
			Usage_BS_Charging
			LogicalMinimum_i8(0)
			LogicalMaximum_i8(1)
			ReportCount(1)
			ReportSize(8)
			Input(Var|Abs)
		)
	)
};

SEC(HID_BPF_RDESC_FIXUP)
int BPF_PROG(ack05_fix_rdesc, struct hid_bpf_ctx *hctx)
{
	__u8 *data = hid_bpf_get_data(hctx, 0 /* offset */, HID_MAX_DESCRIPTOR_SIZE /* size */);
	__s32 rdesc_size = hctx->size;

	if (!data)
		return 0; /* EPERM check */

	if (rdesc_size == VENDOR_DESCRIPTOR_LENGTH) {
		/*
		 * The vendor fixed rdesc is appended after the current one,
		 * to keep the output reports working.
		 */
		__builtin_memcpy(data + rdesc_size, fixed_rdesc_vendor, sizeof(fixed_rdesc_vendor));
		return sizeof(fixed_rdesc_vendor) + rdesc_size;
	}

	hid_set_name(hctx->hid, "Disabled by HID-BPF Hanvon Ugee Shortcut Remote");

	__builtin_memcpy(data, disabled_rdesc, sizeof(disabled_rdesc));
	return sizeof(disabled_rdesc);
}

static int HID_BPF_ASYNC_FUN(switch_to_raw_mode)(struct hid_bpf_ctx *hid)
{
	static __u8 magic_0[32] = {0x02, 0xb0, 0x04, 0x00, 0x00};
	int err;

	/*
	 * The proprietary driver sends the 3 following packets after the
	 * above one.
	 * These don't seem to have any effect, so we don't send them to save
	 * some processing time.
	 *
	 * static __u8 magic_1[32] = {0x02, 0xb4, 0x01, 0x00, 0x01};
	 * static __u8 magic_2[32] = {0x02, 0xb4, 0x01, 0x00, 0xff};
	 * static __u8 magic_3[32] = {0x02, 0xb8, 0x04, 0x00, 0x00};
	 */

	err = hid_bpf_hw_output_report(hid, magic_0, sizeof(magic_0));
	if (err < 0)
		return err;

	return 0;
}

SEC(HID_BPF_DEVICE_EVENT)
int BPF_PROG(ack05_fix_events, struct hid_bpf_ctx *hctx)
{
	__u8 *data = hid_bpf_get_data(hctx, 0 /* offset */, PAD_REPORT_LENGTH);
	int ret = 0;

	if (!data)
		return 0; /* EPERM check */

	if (data[0] != VENDOR_REPORT_ID)
		return 0;

	/* reconnect event */
	if (data[1] == 0xf8 && data[2] == 02 && data[3] == 0x01)
		HID_BPF_ASYNC_DELAYED_CALL(switch_to_raw_mode, hctx, 10);

	/* button event */
	if (data[1] == RAW_PAD_REPORT_ID) {
		data[0] = data[1];
		if (data[7] == 0x02)
			data[7] = 0xff;
		ret = 8;
	} else if (data[1] == RAW_BATTERY_REPORT_ID) {
		data[0] = data[1];
		ret = 5;
	}

	return ret;
}

HID_BPF_OPS(xppen_ack05_remote) = {
	.hid_device_event = (void *)ack05_fix_events,
	.hid_rdesc_fixup = (void *)ack05_fix_rdesc,
};

SEC("syscall")
int probe(struct hid_bpf_probe_args *ctx)
{
	switch (ctx->rdesc_size) {
	case PAD_WIRED_DESCRIPTOR_LENGTH:
	case PAD_DONGLE_DESCRIPTOR_LENGTH:
	case STYLUS_DESCRIPTOR_LENGTH:
	case VENDOR_DESCRIPTOR_LENGTH:
		ctx->retval = 0;
		break;
	default:
		ctx->retval = -EINVAL;
		break;
	}

	if (ctx->rdesc_size == VENDOR_DESCRIPTOR_LENGTH) {
		struct hid_bpf_ctx *hctx = hid_bpf_allocate_context(ctx->hid);

		if (!hctx) {
			ctx->retval = -EINVAL;
			return 0;
		}

		ctx->retval = HID_BPF_ASYNC_INIT(switch_to_raw_mode) ||
			      switch_to_raw_mode(hctx);

		hid_bpf_release_context(hctx);
	}

	return 0;
}

char _license[] SEC("license") = "GPL";
