#include <linux/delay.h>
#include <linux/gpio.h>
#include <linux/i2c.h>
#include <linux/input.h>
#include <linux/interrupt.h>
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/of_gpio.h>

#define DRIVER_NAME	"pmt9123"
#define DEVICE_NAME	"pmt9123"
#define COMPATIBLE_NAME	"pixart,pmt9123"
#define INPUT_NAME	"PMT9123 Optical Track Sensor"

#define DELAY_POLLING_JIFFY		(HZ/60)

#define DELAY_RESET_MIN_US		20
#define DELAY_RESET_MAX_US		20

#define DELAY_INIT_MIN_US		1500
#define DELAY_INIT_MAX_US		2000

#define DELAY_ENABLE_WRITE_MIN_US	300
#define DELAY_ENABLE_WRITE_MAX_US	400

#define DELAY_OBSERVATION_MIN_US	(10 * 1000)
#define DELAY_OBSERVATION_MAX_US	(11 * 1000)

#define I2C_FUNCTIONALITY ( \
	I2C_FUNC_SMBUS_BYTE_DATA | \
	I2C_FUNC_SMBUS_READ_I2C_BLOCK | \
	0)

#define REG_POWER_UP_RESET__FLD_PUR		0xFF
#define REG_POWER_UP_RESET__FLD_PUR__VAL_COLD	0x5A
#define REG_POWER_UP_RESET__FLD_PUR__VAL_WARM	0x96

#define REG_SHUTDOWN__FLD_SD			0xFF
#define REG_SHUTDOWN__FLD_SD__VAL_ENTER		0xE7
#define REG_SHUTDOWN__FLD_SD__VAL_EXIT		0xC3

#define REG_ENABLE_WRITE__FLD_ERW		0xFF
#define REG_ENABLE_WRITE__FLD_ERW__VAL_ENABLE	0xBA
#define REG_ENABLE_WRITE__FLD_ERW__VAL_DISABLE	0xB5

#define REG_RESOLUTION__FLD_RES			0x1F

#define REG_ORIENTATION__FLD_SWAPXY		0x40
#define REG_ORIENTATION__FLD_INVERTX		0x20
#define REG_ORIENTATION__FLD_INVERTY		0x10

#define REG_MOTION__FLD_READY			0x80

#define REG_DELTA_X_L__FLD_X			0xFF

#define REG_DELTA_Y_L__FLD_Y			0xFF

#define REG_DELTA_XY_H__FLD_Y			0x0F
#define REG_DELTA_XY_H__FLD_X			0xF0

#define REG_OBSERVATION_POR			0x1F

#define reg_fld(reg, fld) \
	(REG_##reg##__FLD_##fld)

#define reg_fld_msk(reg, fld) \
	reg_fld(reg, fld)

#define reg_fld_sft(reg, fld) \
	ilog2(reg_fld_msk(reg, fld) ^ (reg_fld_msk(reg, fld) - 1))

#define reg_fld_bts(reg, fld) \
	ilog2(reg_fld_msk(reg, fld) >> reg_fld_sft(reg, fld))

#define reg_fld_get(reg, fld, val) \
	(((val) & reg_fld_msk(reg, fld)) >> reg_fld_sft(reg, fld))

#define reg_fld_set(reg, fld, val) \
	(((val) << reg_fld_sft(reg, fld)) & reg_fld_msk(reg, fld))

#define reg_fld_val(reg, fld, val) \
	reg_fld_set(reg, fld, REG_##reg##__FLD_##fld##__VAL_##val)

enum {
	REG_PRODUCT_ID		= 0x00,
	REG_REVISION_ID		= 0x01,
	REG_MOTION		= 0x02,
	REG_DELTA_X_L		= 0x03,
	REG_DELTA_Y_L		= 0x04,
	REG_DELTA_XY_H		= 0x05,
	REG_SQUAL		= 0x06,
	REG_SHUTTER_U		= 0x07,
	REG_SHUTTER_L		= 0x08,
	REG_PIX_MAX		= 0x09,
	REG_PIX_AVG		= 0x0A,
	REG_PIX_MIN		= 0x0B,
	REG_PERFORMANCE		= 0x11,
	REG_RUN_DOWNSHIFT	= 0x14,
	REG_REST1_RATE		= 0x15,
	REG_REST1_DOWNSHIFT	= 0x16,
	REG_REST2_RATE		= 0x17,
	REG_REST2_DOWNSHIFT	= 0x18,
	REG_REST3_RATE		= 0x19,
	REG_OBSERVATION		= 0x1D,
	REG_FRAME_CAPTURE1	= 0x24,
	REG_FRAME_CAPTURE2	= 0x25,
	REG_POWER_UP_RESET	= 0x3A,
	REG_SHUTDOWN		= 0x3B,
	REG_ENABLE_WRITE	= 0x41,
	REG_RESOLUTION		= 0x48,
	REG_ORIENTATION		= 0x4D,

	REG_AUTO_INCREMENT	= 0x80,
} reg;

struct regval {
	u8 reg;
	u8 val;
};

struct context {
	/* parent */
	struct i2c_client		*client;
	/* children */
	struct input_dev		*input;

	/* data */
	bool				enabled;
	struct {
		struct gpio_desc	*reset;
		struct gpio_desc	*motion;
	}				gpios;
	struct {
		unsigned int		x;
		unsigned int		y;
	}				codes;
	wait_queue_head_t		wait;
	unsigned int			irq;
};

/* Proprietary Optimization Sequence from Datasheet */
static const struct regval perf_configuration[] = {
	/* Tracking Performance */
	{ .reg = 0x11 , .val = 0x0C },
	{ .reg = 0x14 , .val = 0x04 },
	{ .reg = 0x25 , .val = 0x0F },
	{ .reg = 0x27 , .val = 0xAA },
	{ .reg = 0x34 , .val = 0x80 },
	{ .reg = 0x35 , .val = 0x03 },
	{ .reg = 0x39 , .val = 0x89 },
	{ .reg = 0x4D , .val = 0x79 },
	{ .reg = 0x53 , .val = 0x96 },
	{ .reg = 0x58 , .val = 0x66 },
	{ .reg = 0x5D , .val = 0x56 },
	{ .reg = 0x5E , .val = 0xF5 },
	{ .reg = 0x5F , .val = 0xCA },
	{ .reg = 0x61 , .val = 0xD3 },
	{ .reg = 0x6F , .val = 0xEF },
	{ .reg = 0x70 , .val = 0xB2 },
	{ .reg = 0x7D , .val = 0xA2 },
	/* Latency */
	{ .reg = 0x1C , .val = 0x24 },
	{ .reg = 0x26 , .val = 0xA2 },
	{ .reg = 0x56 , .val = 0x32 },
	{ .reg = 0x65 , .val = 0x1F },
	/* Tracking Distance (Z) */
	{ .reg = 0x75 , .val = 0x2A },
	{ .reg = 0x76 , .val = 0x00 },
	{ .reg = 0x77 , .val = 0x12 },
	{ .reg = 0x7A , .val = 0x01 },
	{ .reg = 0x7B , .val = 0x01 },
};

static struct gpio_desc *gpio_init(struct device *device, const char *name,
			    unsigned long flags)
{
	struct gpio_desc *gpio;
	int err;

	gpio = of_get_named_gpiod_flags(device->of_node, name, 0, NULL);
	if (IS_ERR(gpio)) {
		dev_err(device, "%s missing\n", name);
		goto done;
	}

	err = devm_gpio_request_one(device, desc_to_gpio(gpio), flags, NULL);
	if (err < 0) {
		gpio = ERR_PTR(err);
		dev_err(device, "%s request failure\n", name);
		goto done;
	}

done:
	return gpio;
}

static int client_poll(struct i2c_client *client, int *x, int *y)
{
	struct {
		u8 delta_x_l;
		u8 delta_y_l;
		u8 delta_xy_h;
	} data;
	int err;
	struct {
		int x:reg_fld_bts(DELTA_XY_H, X)+reg_fld_bts(DELTA_X_L, X);
		int y:reg_fld_bts(DELTA_XY_H, Y)+reg_fld_bts(DELTA_Y_L, Y);
	} delta;

	err = i2c_smbus_read_i2c_block_data(client,
					    REG_DELTA_X_L | REG_AUTO_INCREMENT,
					    sizeof(data), (u8 *)&data);
	if (err < 0)
		goto done;

	delta.x = reg_fld_get(DELTA_XY_H, X, data.delta_xy_h);
	delta.x = (delta.x << reg_fld_bts(DELTA_X_L, X))
		| reg_fld_get(DELTA_X_L, X, data.delta_x_l);
	delta.y = reg_fld_get(DELTA_XY_H, Y, data.delta_xy_h);
	delta.y = (delta.y << reg_fld_bts(DELTA_Y_L, Y))
		| reg_fld_get(DELTA_Y_L, Y, data.delta_y_l);

	*x += delta.x;
	*y += delta.y;

done:
	return err;
}

static int client_powerup(struct i2c_client *client)
{
	struct context *context = i2c_get_clientdata(client);
	int err;
	size_t i;
	struct {
		u8 motion;
		u8 delta_x_l;
		u8 delta_y_l;
		u8 delta_xy_h;
	} data;

	gpiod_set_value(context->gpios.reset, false);
	usleep_range(DELAY_RESET_MIN_US, DELAY_RESET_MAX_US);
	gpiod_set_value(context->gpios.reset, true);

	usleep_range(DELAY_INIT_MIN_US, DELAY_INIT_MAX_US);

	err = i2c_smbus_write_byte_data(client, REG_ENABLE_WRITE,
					reg_fld_val(ENABLE_WRITE, ERW, ENABLE));
	if (err < 0)
		goto done;

	usleep_range(DELAY_ENABLE_WRITE_MIN_US, DELAY_ENABLE_WRITE_MAX_US);

	err = i2c_smbus_write_byte_data(client, REG_OBSERVATION, 0x00);
	if (err < 0)
		goto done;
	err = i2c_smbus_write_byte_data(client, REG_ENABLE_WRITE,
					reg_fld_val(ENABLE_WRITE,
						    ERW, DISABLE));
	if (err < 0)
		goto done;

	usleep_range(DELAY_OBSERVATION_MIN_US, DELAY_OBSERVATION_MAX_US);

	err = i2c_smbus_read_byte_data(client, REG_OBSERVATION);
	if (err < 0) {
		goto done;
	} else if (err != REG_OBSERVATION_POR) {
		err = -EIO;
		goto done;
	}

	err = i2c_smbus_read_i2c_block_data(client,
					    REG_MOTION | REG_AUTO_INCREMENT,
					    sizeof(data), (u8 *)&data);
	if (err < 0)
		goto done;

	err = i2c_smbus_write_byte_data(client, REG_ENABLE_WRITE,
					reg_fld_val(ENABLE_WRITE, ERW, ENABLE));
	if (err < 0)
		goto done;

	for (i = 0; i < ARRAY_SIZE(perf_configuration); i++) {
		const struct regval *regval = &perf_configuration[i];
		err = i2c_smbus_write_byte_data(client, regval->reg,
						regval->val);
		if (err < 0)
			goto done;
	}

	err = i2c_smbus_write_byte_data(client, REG_ENABLE_WRITE,
					reg_fld_val(ENABLE_WRITE,
						    ERW, DISABLE));
	if (err < 0)
		goto done;

done:
	return err;
}

static int client_powerdown(struct i2c_client *client)
{
	int err;

	err = i2c_smbus_write_byte_data(client, REG_SHUTDOWN,
					reg_fld_val(SHUTDOWN, SD, ENTER));

	return err;
}

static int client_wake(struct i2c_client *client)
{
	int err;

	err = i2c_smbus_write_byte_data(client, REG_SHUTDOWN,
					reg_fld_val(SHUTDOWN, SD, EXIT));

	if (err < 0)
		goto done;

	err = i2c_smbus_write_byte_data(client, REG_POWER_UP_RESET,
					reg_fld_val(POWER_UP_RESET,
						    PUR, WARM));
	if (err < 0)
		goto done;

	usleep_range(DELAY_INIT_MIN_US, DELAY_INIT_MAX_US);

done:
	return err;
}

static int client_sleep(struct i2c_client *client)
{
	int err;

	err = client_powerdown(client);

	return err;
}

static int client_init(struct i2c_client *client, struct context *context)
{
	struct device *device = &client->dev;
	u8 resolution;
	u8 orientation;
	u32 prop;
	bool inverty;
	bool invertx;
	bool swapxy;
	int err;

	i2c_set_clientdata(client, context);

	err = client_powerup(client);
	if (err < 0) {
		dev_err(device, "client power-up failure\n");
		goto done;
	}

	err = i2c_smbus_write_byte_data(client, REG_ENABLE_WRITE,
					reg_fld_val(ENABLE_WRITE, ERW, ENABLE));
	if (err < 0)
		goto done;

	if (!of_property_read_u32(device->of_node, "resolution", &prop)) {

		resolution = reg_fld_set(RESOLUTION, RES, prop);

		err = i2c_smbus_write_byte_data(client, REG_RESOLUTION,
						resolution);
		if (err < 0)
			goto done;
	}

	inverty = of_property_read_bool(device->of_node, "y-invert");
	invertx = of_property_read_bool(device->of_node, "x-invert");
	swapxy = of_property_read_bool(device->of_node, "x-y-swap");

	orientation =
		reg_fld_set(ORIENTATION, INVERTY, inverty) |
		reg_fld_set(ORIENTATION, INVERTX, invertx) |
		reg_fld_set(ORIENTATION, SWAPXY, swapxy);

	err = i2c_smbus_write_byte_data(client, REG_ORIENTATION, orientation);
	if (err < 0)
		goto done;

	err = i2c_smbus_write_byte_data(client, REG_ENABLE_WRITE,
					reg_fld_val(ENABLE_WRITE,
						    ERW, DISABLE));
	if (err < 0)
		goto done;

done:
	return err;
}

static int client_deinit(struct i2c_client *client)
{
	int err;

	err = client_powerdown(client);

	return err;
}

static void input_report(struct input_dev *input, int x, int y)
{
	struct context *context = input_get_drvdata(input);

	input_event(input, EV_REL, context->codes.x, x);
	input_event(input, EV_REL, context->codes.y, y);
	input_sync(input);
}

static int input_open(struct input_dev *input)
{
	struct context *context = input_get_drvdata(input);
	struct i2c_client *client = context->client;
	int err;

	err = client_wake(client);
	if (err < 0)
		goto done;

	context->enabled = true;
	mb();

	enable_irq(context->irq);

done:
	return err;
}

static void input_close(struct input_dev *input)
{
	struct context *context = input_get_drvdata(input);
	struct i2c_client *client = context->client;

	context->enabled = false;
	mb();
	wake_up(&context->wait);

	disable_irq(context->irq);

	client_sleep(client);
}

static struct input_dev *input_init(struct device *device,
				    struct context *context)
{
	struct input_dev *input;
	bool valid = false;
	unsigned int code;
	int err;

	input = devm_input_allocate_device(device);
	if (!input) {
		input = ERR_PTR(-ENOMEM);
		dev_err(device, "input allocation failure\n");
		goto done;
	}

	input_set_drvdata(input, context);

	input->name = INPUT_NAME;
	input->id.bustype = BUS_I2C;
	__set_bit(EV_REL, input->evbit);
	input->open = input_open;
	input->close = input_close;

	err = of_property_read_u32(device->of_node, "x-input-code", &code);
	if (err < 0)
		code = REL_CNT;
	if (code < REL_CNT) {
		__set_bit(code, input->relbit);
		valid = true;
	}
	context->codes.x = code;

	err = of_property_read_u32(device->of_node, "y-input-code", &code);
	if (err < 0)
		code = REL_CNT;
	if (code < REL_CNT) {
		__set_bit(code, input->relbit);
		valid = true;
	}
	context->codes.y = code;

	if (!valid) {
		input = ERR_PTR(-EINVAL);
		goto done;
	}

	err = input_register_device(input);
	if (err < 0) {
		input = ERR_PTR(err);
		dev_err(device, "input registration failure\n");
		goto done;
	}

done:
	return input;
}

static irqreturn_t irq_thread(int irq, void *handle)
{
	struct context *context = handle;
	struct i2c_client *client = context->client;
	struct device *device = &client->dev;
	struct input_dev *input = context->input;
	int err;

	while (context->enabled &&
	       !gpiod_get_value(context->gpios.motion)) {

		int x = 0;
		int y = 0;

		err = client_poll(client, &x, &y);
		if (err < 0)
			dev_err(device, "client poll failure\n");

		input_report(input, x, y);

		wait_event_timeout(context->wait, !context->enabled,
				   DELAY_POLLING_JIFFY);
	}

	return IRQ_HANDLED;
}

static int irq_init(struct device *device, struct context *context,
		    wait_queue_head_t *wait, struct gpio_desc *gpio)
{
	int irq;
	int err;

	init_waitqueue_head(wait);

	irq = gpiod_to_irq(gpio);
	if (irq < 0)
		goto done;

	err = gpiod_lock_as_irq(gpio);
	if (err < 0) {
		irq = err;
		goto done;
	}

	err = devm_request_threaded_irq(device, irq, NULL, irq_thread,
					IRQF_TRIGGER_FALLING | IRQF_ONESHOT,
					DRIVER_NAME, context);
	if (err < 0) {
		irq = err;
		dev_err(device, "irq request failure\n");
		goto done;
	}

	disable_irq(irq);

done:
	return irq;
}

static int pmt9123_probe(struct i2c_client *client,
			 const struct i2c_device_id *id)
{
	struct i2c_adapter *adapter = client->adapter;
	struct device *device = &client->dev;
	struct context *context;
	struct gpio_desc *gpio;
	struct input_dev *input;
	int irq;
	int err;

	if (!i2c_check_functionality(adapter, I2C_FUNCTIONALITY)) {
		err = -EIO;
		dev_err(device, "i2c functionality failure\n");
		goto done;
	}

	if (!device->of_node) {
		err = -ENODEV;
		dev_err(device, "device tree node missing\n");
		goto done;
	}

	context = devm_kzalloc(device, sizeof(*context), GFP_KERNEL);
	if (!context) {
		err = -ENOMEM;
		dev_err(device, "context allocation failure\n");
		goto done;
	}

	gpio = gpio_init(device, "reset-gpios",
			 GPIOF_OUT_INIT_HIGH |
			 GPIOF_EXPORT_DIR_FIXED);
	if (IS_ERR(gpio)) {
		err = PTR_ERR(gpio);
		goto done;
	}
	context->gpios.reset = gpio;

	gpio = gpio_init(device, "motion-gpios",
			 GPIOF_IN |
			 GPIOF_EXPORT_DIR_FIXED);
	if (IS_ERR(gpio)) {
		err = PTR_ERR(gpio);
		goto done;
	}
	context->gpios.motion = gpio;

	err = client_init(client, context);
	if (err < 0)
		goto done;
	context->client = client;

	irq = irq_init(device, context, &context->wait, context->gpios.motion);
	if (irq < 0) {
		err = irq;
		goto done;
	}
	context->irq = irq;

	input = input_init(device, context);
	if (IS_ERR(input)) {
		err = PTR_ERR(input);
		goto done;
	}
	context->input = input;

	device_init_wakeup(device, true);

	client_sleep(client);

done:
	return err;
}

static int pmt9123_remove(struct i2c_client *client)
{
	struct device *device = &client->dev;
	int err;

	device_init_wakeup(device, false);

	err = client_deinit(client);

	return err;
}

#ifdef CONFIG_PM_SLEEP

static int pmt9123_suspend(struct device *device)
{
	struct i2c_client *client = to_i2c_client(device);
	struct context *context = i2c_get_clientdata(client);

	if (device_may_wakeup(device))
		enable_irq_wake(context->irq);
	else if (context->enabled)
		client_sleep(client);

	return 0;
}

static int pmt9123_resume(struct device *device)
{
	struct i2c_client *client = to_i2c_client(device);
	struct context *context = i2c_get_clientdata(client);
	int err = 0;

	if (device_may_wakeup(device))
		disable_irq_wake(context->irq);
	else if (context->enabled)
		err = client_wake(client);

	return err;
}

#endif /* CONFIG_PM_SLEEP */

static SIMPLE_DEV_PM_OPS(pmt9123_pm_ops, pmt9123_suspend, pmt9123_resume);

static struct i2c_device_id pmt9123_idtable[] = {
	{ .name = DEVICE_NAME },
	{ }
};

MODULE_DEVICE_TABLE(i2c, pmt9123_idtable);

static const struct of_device_id pmt9123_oftable[] = {
	{ .compatible = COMPATIBLE_NAME },
	{ }
};
MODULE_DEVICE_TABLE(of, pmt9123_oftable);

static struct i2c_driver pmt9123_driver = {
	.driver.name		= DRIVER_NAME,
	.driver.of_match_table	= of_match_ptr(pmt9123_oftable),
	.driver.pm		= &pmt9123_pm_ops,

	.id_table		= pmt9123_idtable,
	.probe			= pmt9123_probe,
	.remove			= pmt9123_remove,
};

module_i2c_driver(pmt9123_driver);

MODULE_AUTHOR("Nestlabs, Inc.");
MODULE_DESCRIPTION("Driver for Pixart PMT9123 Optical Track Sensor");
MODULE_LICENSE("GPLv2");