#include <linux/module.h>
#include <linux/slab.h>
#include <linux/hid.h>
#include <linux/power_supply.h>
#include <linux/init.h>
#include <linux/netlink.h>
#include <linux/skbuff.h>
#include <net/sock.h>
#include <linux/kernel.h>

#define JBL_VID  0x0ecb
#define JBL_PID  0x205c
#define MY_NETLINK_GROUP 27
//#define NETLINK_MYPROTO 31

static struct sock *nl_sk = NULL;

/* List of HID devices this driver supports */
static const struct hid_device_id jbl_hid_table[] = {
    { HID_USB_DEVICE(JBL_VID, JBL_PID) },
    { }
};
MODULE_DEVICE_TABLE(hid, jbl_hid_table);


static void jbl_netlink_broadcast(const char *msg) {
    struct sk_buff *skb_out;
    struct nlmsghdr *nlh;
    int msg_size = strlen(msg);
    int res;

    if (!nl_sk) {
        pr_err("jbl610: netlink socket not ready!\n");
        return;
    }

    skb_out = nlmsg_new(msg_size, GFP_KERNEL);
    if (!skb_out) {
        pr_err("jbl610: Failed to allocate skb\n");
        return;
    }

    nlh = nlmsg_put(skb_out, 0, 0, NLMSG_DONE, msg_size, 0);
    if (!nlh) {
        pr_err("jbl610: nlmsg_put() failed\n");
        kfree_skb(skb_out);
        return;
    }

    memcpy(nlmsg_data(nlh), msg, msg_size);
    res = netlink_broadcast(nl_sk, skb_out, 0, MY_NETLINK_GROUP, GFP_KERNEL);
    //res = netlink_unicast(nl_sk, skb_out, 0, MY_NETLINK_GROUP, GFP_KERNEL);
    if (res < 0) {
        pr_err("jbl610: netlink_broadcast error: %d\n", res);
    } else {
        pr_info("jbl610: broadcasted message: \"%s\"\n", msg);
    }
}

struct jbl_headset_data {
    struct hid_device *hdev;
    struct power_supply *battery;

    int capacity;

    int status;  
    int present;
    int muted;
};

static enum power_supply_property jbl_battery_props[] = {
    POWER_SUPPLY_PROP_CAPACITY,
    POWER_SUPPLY_PROP_STATUS,
    POWER_SUPPLY_PROP_MODEL_NAME,
    POWER_SUPPLY_PROP_MANUFACTURER,
    POWER_SUPPLY_PROP_PRESENT
};

static int jbl_battery_get_property(
    struct power_supply *psy,
    enum power_supply_property psp,
    union power_supply_propval *val)
{
    struct jbl_headset_data *jbl = power_supply_get_drvdata(psy);

    switch (psp) {
    case POWER_SUPPLY_PROP_CAPACITY:
        val->intval = jbl->capacity;  // e.g. 0..100
        return 0;
    case POWER_SUPPLY_PROP_STATUS:

        val->intval = jbl->status;  
        return 0;
     case POWER_SUPPLY_PROP_MODEL_NAME:
        val->strval = "JBL Quantum 610";
        return 0;
    case POWER_SUPPLY_PROP_MANUFACTURER:
        val->strval = "JBL";
        return 0;
    case POWER_SUPPLY_PROP_PRESENT:
        val->intval = jbl->present;
        return 0;
    default:
        return -EINVAL;
    }
}

static const struct power_supply_desc jbl_battery_desc = {
    .name = "JBL_Quantum_610",
    .type = POWER_SUPPLY_TYPE_WIRELESS,
    .properties = jbl_battery_props,
    .num_properties = ARRAY_SIZE(jbl_battery_props),
    .get_property = jbl_battery_get_property,
};

static int jbl_raw_event(struct hid_device *hdev, struct hid_report *report,
                         u8 *data, int size)
{
    struct jbl_headset_data *jbl = hid_get_drvdata(hdev);

    if (size > 1 && data[0] == 0x08) {

        jbl->present = true; // if we have any reports from device, that's mean that it is now connected
        int new_cap = data[1]; 
        if (new_cap != jbl->capacity) {
            dev_info(&hdev->dev, "Capacity changed, old value: %d, new: %d", jbl->capacity, new_cap);
            if(new_cap > jbl->capacity) {
                jbl->status = POWER_SUPPLY_STATUS_CHARGING;
            } else {
                jbl->status = POWER_SUPPLY_STATUS_DISCHARGING;
            }


            if(new_cap == 100) {
                jbl->status = POWER_SUPPLY_STATUS_FULL;
            }
            char buffer[30];
            sprintf(buffer, "JBL610_BATTERY=%d", new_cap);
            jbl_netlink_broadcast(buffer);

            jbl->capacity = new_cap;
            power_supply_changed(jbl->battery);
        }
    }

    // device disconnected
    if(size > 1 && data[0] == 0x09) {
        jbl->present = 0;
    }

    // device connected 
    if(size > 1 && data[0] == 0x07 && data[1] == 0x01) {
        jbl->present = 1;
    }

    if(size > 1 && data[0] == 0x06) {
        int muted = data[1] == 0 ? 1 : 0;
        char buf[20];
        if(muted) {
            sprintf(buf, "JBL610_DEVICE_MUTED");
        } else {
            sprintf(buf, "JBL610_DEVICE_UNMUTED");
        }
        jbl_netlink_broadcast(buf);
        
        jbl->muted = muted;
    }

    return 0;
}

static ssize_t mic_mute_show(struct device *dev,
                             struct device_attribute *attr,
                             char *buf)
{
    struct jbl_headset_data *jbl = dev_get_drvdata(dev);
    return sprintf(buf, "%d\n", jbl->muted ? 1 : 0);
}


static DEVICE_ATTR_RO(mic_mute); // read-only attribute, for example

static int jbl_probe(struct hid_device *hdev, const struct hid_device_id *id)
{
    int ret;
    struct jbl_headset_data *jbl;
    struct power_supply_config psy_cfg = { };

    /* Allocate memory for our driver data */
    jbl = devm_kzalloc(&hdev->dev, sizeof(*jbl), GFP_KERNEL);
    if (!jbl)
        return -ENOMEM;

    jbl->hdev = hdev;
    hid_set_drvdata(hdev, jbl);

    /* Register with the HID subsystem */
    ret = hid_parse(hdev);
    if (ret) {
        dev_err(&hdev->dev, "hid_parse failed\n");
        return ret;
    }

    ret = hid_hw_start(hdev, HID_CONNECT_DEFAULT);
    if (ret) {
        dev_err(&hdev->dev, "hid_hw_start failed\n");
        return ret;
    }

    /* Setup power_supply */
    psy_cfg.drv_data = jbl;  // we can retrieve jbl in get_property
    jbl->capacity = 50;      // default (unknown) until we get a real reading
    jbl->status   = POWER_SUPPLY_STATUS_DISCHARGING;

    jbl->battery = devm_power_supply_register(&hdev->dev,
                                              &jbl_battery_desc,
                                              &psy_cfg);
    if (IS_ERR(jbl->battery)) {
        dev_err(&hdev->dev, "Failed to register power supply\n");
        hid_hw_stop(hdev);
        return PTR_ERR(jbl->battery);
    }

    ret = device_create_file(&hdev->dev, &dev_attr_mic_mute);
    if (ret) {
        dev_err(&hdev->dev, "Failed to create mic_mute attribute\n");
        return ret;
    }

    struct netlink_kernel_cfg cfg = {
        .input = NULL,  // We do NOT receive from user space
        .groups = 0,
    };

    pr_info("jbl610: netlink Initializing...\n");

    nl_sk = netlink_kernel_create(&init_net, NETLINK_USERSOCK, &cfg);
    //nl_sk =netlink_kernel_create(&init_net, NETLINK_USERSOCK, 1, jbl_nl_receive_callback, NULL, THIS_MODULE)
    if (!nl_sk) {
        pr_err("jbl610: Error creating netlink socket.\n");
        return -ENOMEM;
    }


    /* Example: Send a test broadcast right away (for demonstration). */
   // jbl_netlink_broadcast("jbl_driver_initialized");


    dev_info(&hdev->dev, "JBL headset driver probed\n");
    return 0;
}

static void jbl_nl_receive_callback(struct sk_buff *skb) {
    nlmsg_free(skb);
}

/* Called when the driver is unbound or device is removed */
static void jbl_remove(struct hid_device *hdev)
{
    if (nl_sk) {
        netlink_kernel_release(nl_sk);
        nl_sk = NULL;
    }
    pr_info("jbl610: Exiting.\n");
    device_remove_file(&hdev->dev, &dev_attr_mic_mute);
    hid_hw_stop(hdev); 
    /* power_supply is auto-freed by devm_ once the device is gone */
}

static struct hid_driver jbl_driver = {
    .name       = "jbl_headset",
    .id_table   = jbl_hid_table,
    .probe      = jbl_probe,
    .remove     = jbl_remove,
    .raw_event  = jbl_raw_event,  // We'll get HID input packets here
};

module_hid_driver(jbl_driver);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Kirill Ivlev <webster.spb@gmail.com>");
MODULE_DESCRIPTION("Kernel driver for JBL Headset battery");