#include <linux/kernel.h>
#include <linux/device.h>
#include <linux/hid.h>
#include <linux/module.h>
#include <linux/usb.h>
#include <linux/power_supply.h>

#define USB_VID 0x258A
#define USB_PID 0x002F
#define USB_PID_WIRED 0x002E

#define USB_BATTERY_TIMEOUT_MS 15000

struct redragon_m811_device {
    struct hid_device *hdev;
    bool removed;
    bool workable;

    struct delayed_work battery_work;

    struct power_supply_desc battery_desc;
    struct power_supply *battery;
    uint8_t battery_capacity;
    bool connected;
    bool wireless;
    bool idle;
};

static enum power_supply_property redragon_m811_battery_props[] = {
    POWER_SUPPLY_PROP_PRESENT,
	POWER_SUPPLY_PROP_STATUS,
	POWER_SUPPLY_PROP_SCOPE,
	POWER_SUPPLY_PROP_CAPACITY,
};

static int redragon_m811_battery_get_property(struct power_supply *psy,
                                     enum power_supply_property psp,
                                     union power_supply_propval *val) {
    struct redragon_m811_device *m811 = power_supply_get_drvdata(psy);
    int ret = 0;

    switch (psp) {
    case POWER_SUPPLY_PROP_PRESENT:
		val->intval = 1;
		break;
    case POWER_SUPPLY_PROP_STATUS:
        if (!m811)
            return -EINVAL;
		val->intval = m811->connected ?
			(m811->wireless ? POWER_SUPPLY_STATUS_DISCHARGING : POWER_SUPPLY_STATUS_CHARGING) :
			POWER_SUPPLY_STATUS_UNKNOWN;
        if (m811->battery_capacity == 100 && !m811->wireless)
            val->intval = POWER_SUPPLY_STATUS_FULL;
		break;
	case POWER_SUPPLY_PROP_SCOPE:
		val->intval = POWER_SUPPLY_SCOPE_DEVICE;
		break;
	case POWER_SUPPLY_PROP_CAPACITY:
        if (!m811)
            return -EINVAL;
		val->intval = m811->battery_capacity;
		break;
    default:
        ret = -EINVAL;
        break;
    }

    return ret;
}

static int __get_data(struct hid_device *hdev, u8 *buf, unsigned char dat) {
    int ret;

    memset(buf, 0, sizeof(unsigned char) * 8);
    buf[0] = 0x05;
    buf[1] = dat;

    ret = hid_hw_raw_request(hdev, 0x05,
				 buf, sizeof(unsigned char) * 8,
				 0x02, HID_REQ_SET_REPORT);
	if (ret < (int)(sizeof(unsigned char) * 8)) {
		hid_err(hdev, "hid_hw_raw_request() failed with %d\n", ret);
		ret = -ENODATA;
	}

    memset(buf, 0, sizeof(unsigned char) * 8);
    buf[0] = 0x05;
    buf[1] = dat;

	ret = hid_hw_raw_request(hdev, 0x05,
				 buf, sizeof(unsigned char) * 8,
				 0x02, HID_REQ_GET_REPORT);
	if (ret < (int)(sizeof(unsigned char) * 8)) {
		hid_err(hdev, "hid_hw_raw_request() failed with %d\n", ret);
		ret = -ENODATA;
	}

    printk(KERN_INFO "GET_REQUEST(0x%02x): %02x %02x %02x %02x", dat, buf[0], buf[1], buf[2], buf[3]);

    return ret;
}

static int redragon_m811_fetch_battery(struct hid_device *hdev)
{
	struct redragon_m811_device *m811 = hid_get_drvdata(hdev);
	int ret = 0;
    u8 *write_buf;

    if (!m811->workable) return ret;

    printk(KERN_INFO "Fetchig battery data");

	write_buf = kmalloc(sizeof(unsigned char) * 8, GFP_KERNEL);
	if (!write_buf) return -ENOMEM;

    ret = __get_data(hdev, write_buf, 0x80); // idle status
    if(ret < 0)
        return ret;

    m811->idle = false;
    if (write_buf[1] == 0x80 && write_buf[2] == 0x00) // idle
        m811->idle = true;

    if(!m811->idle) {
        ret = __get_data(hdev, write_buf, 0x90); // battery status
        if(ret < 0)
            return ret;

        m811->wireless = true;
        m811->battery_capacity = write_buf[3];
        if (m811->battery_capacity > 100)
            m811->battery_capacity = 69;

        if (write_buf[2] == 0x10) {
            m811->wireless = false;
            m811->battery_capacity = 69;
            if (write_buf[3] == 0x02)
                m811->battery_capacity = 100;
        }
    }
    
	kfree(write_buf);

    return ret;
}

static void redragon_m811_battery_timer_tick(struct work_struct *work)
{
	struct redragon_m811_device *m811 = container_of(work,
		struct redragon_m811_device, battery_work.work);
	struct hid_device *hdev = m811->hdev;

    if(m811 && !m811->removed && m811->workable) {
        redragon_m811_fetch_battery(hdev);

        schedule_delayed_work(&m811->battery_work,
                    msecs_to_jiffies(USB_BATTERY_TIMEOUT_MS));
    }
}

static atomic_t battery_no = ATOMIC_INIT(0);
static int redragon_m811_battery_register(struct redragon_m811_device *m811) {
    struct power_supply_config battery_cfg = { .drv_data = m811, };
	unsigned long n;
	int ret;

	n = atomic_inc_return(&battery_no) - 1;
	m811->battery_desc.name = devm_kasprintf(&m811->hdev->dev, GFP_KERNEL,
						    "redragon_m811_battery_%ld", n);
    if (!m811->battery_desc.name)
		return -ENOMEM;
    m811->battery_desc.type = POWER_SUPPLY_TYPE_BATTERY;
    m811->battery_desc.get_property = redragon_m811_battery_get_property;
    m811->battery_desc.properties = redragon_m811_battery_props;
    m811->battery_desc.num_properties = ARRAY_SIZE(redragon_m811_battery_props);
    m811->battery_desc.use_for_apm = 0;
	m811->battery_capacity = 100;

    m811->battery = devm_power_supply_register(&m811->hdev->dev,
			&m811->battery_desc, &battery_cfg);
    if (IS_ERR(m811->battery)) {
		ret = PTR_ERR(m811->battery);
		hid_err(m811->hdev,
				"%s:power_supply_register failed with error %d\n",
				__func__, ret);
		return ret;
	}
    power_supply_powers(m811->battery, &m811->hdev->dev);

    printk(KERN_INFO "Added power_supply");

    INIT_DELAYED_WORK(&m811->battery_work, redragon_m811_battery_timer_tick);
	redragon_m811_fetch_battery(m811->hdev);
    schedule_delayed_work(&m811->battery_work,
				msecs_to_jiffies(USB_BATTERY_TIMEOUT_MS));
    return 0;
}

static int redragon_m811_device_probe(struct hid_device *hdev,
                            const struct hid_device_id *id) {
    struct redragon_m811_device *m811;
    int ret;

    m811 = devm_kzalloc(&hdev->dev, sizeof(*m811), GFP_KERNEL);
	if (!m811)
		return -ENOMEM;
    
    hid_set_drvdata(hdev, m811);
    ret = hid_parse(hdev);
	if (ret)
		return ret;
    ret = hid_hw_start(hdev, HID_CONNECT_DEFAULT);
	if (ret)
		return ret;

    m811->hdev = hdev;
    m811->wireless = id->driver_data;
    m811->connected = true;

    printk(KERN_INFO "Add device");
    if(hdev->type == HID_TYPE_OTHER) {
        printk(KERN_INFO "Found device Redragon M811 mouse");
        m811->workable = true;

        if (redragon_m811_battery_register(m811) < 0)
            hid_err(m811->hdev, "Failed to register battery for Redragon M811 mouse\n");
    }

    return 0;
}

static void redragon_m811_device_remove(struct hid_device *hdev) {
    struct redragon_m811_device *m811 = hid_get_drvdata(hdev);
    
    if (m811) {
        if (m811->workable) {
            cancel_delayed_work_sync(&m811->battery_work);
            printk(KERN_INFO "Cancelled delayed work");
        }
        m811->removed = true;
    }
    printk(KERN_INFO "Removed device");
    hid_hw_stop(hdev);
}

static const struct hid_device_id redragon_m811_device_id_table[] = {
    { HID_USB_DEVICE(USB_VID, USB_PID), .driver_data = true },
    { HID_USB_DEVICE(USB_VID, USB_PID_WIRED), .driver_data = false },
    {}
};

MODULE_DEVICE_TABLE(hid, redragon_m811_device_id_table);

static struct hid_driver redragon_m811_battery_driver = {
    .name = "redragon_m811_battery",
    .id_table = redragon_m811_device_id_table,
    .probe = redragon_m811_device_probe,
    .remove = redragon_m811_device_remove,
};

module_hid_driver(redragon_m811_battery_driver);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Tobias Manherz");
MODULE_DESCRIPTION("Redragon M811 battery driver");