// SPDX-License-Identifier: GPL-2.0-or-later
/*
 *  Osram AMS AS3668 LED Driver IC
 *
 *  Copyright (C) 2025 Lukas Timmermann <linux@timmermann.space>
 */

#include <linux/bitfield.h>
#include <linux/i2c.h>
#include <linux/leds.h>
#include <linux/module.h>
#include <linux/uleds.h>

#define AS3668_MAX_LEDS                 4

/* Chip Ident */

#define AS3668_CHIP_ID1_REG             0x3e
#define AS3668_CHIP_ID                  0xa5

/* Current Control */

#define AS3668_CURR_MODE_REG            0x01
#define AS3668_CURR_MODE_OFF            0x0
#define AS3668_CURR_MODE_ON             0x1
#define AS3668_CURR1_MODE_MASK          GENMASK(1, 0)
#define AS3668_CURR2_MODE_MASK          GENMASK(3, 2)
#define AS3668_CURR3_MODE_MASK          GENMASK(5, 4)
#define AS3668_CURR4_MODE_MASK          GENMASK(7, 6)
#define AS3668_CURR1_REG                0x02
#define AS3668_CURR2_REG                0x03
#define AS3668_CURR3_REG                0x04
#define AS3668_CURR4_REG                0x05

#define AS3668_CURR_MODE_PACK(mode)     (((mode) << 0) | \
                                        ((mode) << 2) | \
                                        ((mode) << 4) | \
                                        ((mode) << 6))

struct as3668_led {
        struct led_classdev cdev;
        struct as3668 *chip;
        struct fwnode_handle *fwnode;
        u8 mode_mask;
        u8 current_reg;
};

struct as3668 {
        struct i2c_client *client;
        struct as3668_led leds[AS3668_MAX_LEDS];
};

static int as3668_channel_mode_set(struct as3668_led *led, u8 mode)
{
        int ret;
        u8 channel_modes;

        ret = i2c_smbus_read_byte_data(led->chip->client, AS3668_CURR_MODE_REG);
        if (ret < 0) {
                dev_err(led->cdev.dev, "failed to read channel modes\n");
                return ret;
        }
        channel_modes = (u8)ret;

        channel_modes &= ~led->mode_mask;
        channel_modes |= led->mode_mask & (AS3668_CURR_MODE_PACK(mode));

        return i2c_smbus_write_byte_data(led->chip->client, AS3668_CURR_MODE_REG, channel_modes);
}

static enum led_brightness as3668_brightness_get(struct led_classdev *cdev)
{
        struct as3668_led *led = container_of(cdev, struct as3668_led, cdev);

        return i2c_smbus_read_byte_data(led->chip->client, led->current_reg);
}

static void as3668_brightness_set(struct led_classdev *cdev, enum led_brightness brightness)
{
        struct as3668_led *led = container_of(cdev, struct as3668_led, cdev);
        int err;

        err = as3668_channel_mode_set(led, !!brightness);
        if (err)
                dev_err(cdev->dev, "failed to set channel mode: %d\n", err);

        err = i2c_smbus_write_byte_data(led->chip->client, led->current_reg, brightness);
        if (err)
                dev_err(cdev->dev, "failed to set brightness: %d\n", err);
}

static int as3668_dt_init(struct as3668 *as3668)
{
        struct device *dev = &as3668->client->dev;
        struct as3668_led *led;
        struct led_init_data init_data = {};
        int err;
        u32 reg;

        for_each_available_child_of_node_scoped(dev_of_node(dev), child) {
                err = of_property_read_u32(child, "reg", &reg);
                if (err)
                        return dev_err_probe(dev, err, "failed to read 'reg' property");

                if (reg < 0 || reg >= AS3668_MAX_LEDS)
                        return dev_err_probe(dev, -EINVAL,
                                             "unsupported LED: %d\n", reg);

                led = &as3668->leds[reg];
                led->fwnode = of_fwnode_handle(child);

                led->current_reg = reg + AS3668_CURR1_REG;
                led->mode_mask = AS3668_CURR1_MODE_MASK << (reg * 2);
                led->chip = as3668;

                led->cdev.max_brightness = U8_MAX;
                led->cdev.brightness_get = as3668_brightness_get;
                led->cdev.brightness_set = as3668_brightness_set;

                init_data.fwnode = led->fwnode;
                init_data.default_label = ":";

                err = devm_led_classdev_register_ext(dev, &led->cdev, &init_data);
                if (err)
                        return dev_err_probe(dev, err, "failed to register LED %d\n", reg);
        }

        return 0;
}

static int as3668_probe(struct i2c_client *client)
{
        struct as3668 *as3668;
        int err;
        u8 chip_id;

        chip_id = i2c_smbus_read_byte_data(client, AS3668_CHIP_ID1_REG);
        if (chip_id != AS3668_CHIP_ID)
                return dev_err_probe(&client->dev, -ENODEV,
                                     "expected chip ID 0x%02x, got 0x%02x\n",
                                     AS3668_CHIP_ID, chip_id);

        as3668 = devm_kzalloc(&client->dev, sizeof(*as3668), GFP_KERNEL);
        if (!as3668)
                return -ENOMEM;

        as3668->client = client;

        err = as3668_dt_init(as3668);
        if (err)
                return err;

        /* Set all four channel modes to 'off' */
        err = i2c_smbus_write_byte_data(client, AS3668_CURR_MODE_REG,
                                        FIELD_PREP(AS3668_CURR1_MODE_MASK, AS3668_CURR_MODE_OFF) |
                                        FIELD_PREP(AS3668_CURR2_MODE_MASK, AS3668_CURR_MODE_OFF) |
                                        FIELD_PREP(AS3668_CURR3_MODE_MASK, AS3668_CURR_MODE_OFF) |
                                        FIELD_PREP(AS3668_CURR4_MODE_MASK, AS3668_CURR_MODE_OFF));

        /* Set initial currents to 0mA */
        err |= i2c_smbus_write_byte_data(client, AS3668_CURR1_REG, 0);
        err |= i2c_smbus_write_byte_data(client, AS3668_CURR2_REG, 0);
        err |= i2c_smbus_write_byte_data(client, AS3668_CURR3_REG, 0);
        err |= i2c_smbus_write_byte_data(client, AS3668_CURR4_REG, 0);

        if (err)
                return dev_err_probe(&client->dev, -EIO, "failed to set zero initial current levels\n");

        return 0;
}

static void as3668_remove(struct i2c_client *client)
{
        i2c_smbus_write_byte_data(client, AS3668_CURR_MODE_REG, 0);
}

static const struct i2c_device_id as3668_idtable[] = {
        { "as3668" },
        { }
};
MODULE_DEVICE_TABLE(i2c, as3668_idtable);

static const struct of_device_id as3668_match_table[] = {
        { .compatible = "ams,as3668" },
        { }
};
MODULE_DEVICE_TABLE(of, as3668_match_table);

static struct i2c_driver as3668_driver = {
        .driver = {
                .name = "leds_as3668",
                .of_match_table = as3668_match_table,
        },
        .probe = as3668_probe,
        .remove = as3668_remove,
        .id_table = as3668_idtable,
};
module_i2c_driver(as3668_driver);

MODULE_AUTHOR("Lukas Timmermann <linux@timmermann.space>");
MODULE_DESCRIPTION("AS3668 LED driver");
MODULE_LICENSE("GPL");