/*
 * Copyright (c) 2023 Antmicro <www.antmicro.com>
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#define DT_DRV_COMPAT ambiq_i2c

#include <errno.h>
#include <zephyr/drivers/i2c.h>
#include <zephyr/kernel.h>
#include <zephyr/pm/device.h>
#include <zephyr/pm/policy.h>
#include <zephyr/pm/device_runtime.h>

#include <am_mcu_apollo.h>

#ifdef CONFIG_I2C_AMBIQ_BUS_RECOVERY
#include <zephyr/drivers/gpio.h>
#include "i2c_bitbang.h"
#endif /* CONFIG_I2C_AMBIQ_BUS_RECOVERY */

#include <zephyr/logging/log.h>
#include <zephyr/drivers/pinctrl.h>

LOG_MODULE_REGISTER(ambiq_i2c, CONFIG_I2C_LOG_LEVEL);

typedef int (*ambiq_i2c_pwr_func_t)(void);

#define PWRCTRL_MAX_WAIT_US       5
#define I2C_TRANSFER_TIMEOUT_MSEC 500 /* Transfer timeout period */

#include "i2c-priv.h"

struct i2c_ambiq_config {
#ifdef CONFIG_I2C_AMBIQ_BUS_RECOVERY
	struct gpio_dt_spec scl;
	struct gpio_dt_spec sda;
#endif /* CONFIG_I2C_AMBIQ_BUS_RECOVERY */
	uint32_t base;
	int size;
	uint32_t bitrate;
	const struct pinctrl_dev_config *pcfg;
	ambiq_i2c_pwr_func_t pwr_func;
	void (*irq_config_func)(void);
};

typedef void (*i2c_ambiq_callback_t)(const struct device *dev, int result, void *data);

struct i2c_ambiq_data {
	am_hal_iom_config_t iom_cfg;
	void *iom_handler;
	struct k_sem bus_sem;
	struct k_sem transfer_sem;
	i2c_ambiq_callback_t callback;
	void *callback_data;
	int inst_idx;
	uint32_t transfer_status;
	bool pm_policy_state_on;
};

static void i2c_ambiq_pm_policy_state_lock_get(const struct device *dev)
{
	if (IS_ENABLED(CONFIG_PM)) {
		struct i2c_ambiq_data *data = dev->data;

		if (!data->pm_policy_state_on) {
			data->pm_policy_state_on = true;
			pm_policy_state_lock_get(PM_STATE_SUSPEND_TO_RAM, PM_ALL_SUBSTATES);
			pm_device_runtime_get(dev);
		}
	}
}

static void i2c_ambiq_pm_policy_state_lock_put(const struct device *dev)
{
	if (IS_ENABLED(CONFIG_PM)) {
		struct i2c_ambiq_data *data = dev->data;

		if (data->pm_policy_state_on) {
			data->pm_policy_state_on = false;
			pm_device_runtime_put(dev);
			pm_policy_state_lock_put(PM_STATE_SUSPEND_TO_RAM, PM_ALL_SUBSTATES);
		}
	}
}

#ifdef CONFIG_I2C_AMBIQ_DMA
static __aligned(32) struct {
	__aligned(32) uint32_t buf[CONFIG_I2C_DMA_TCB_BUFFER_SIZE];
} i2c_dma_tcb_buf[DT_NUM_INST_STATUS_OKAY(DT_DRV_COMPAT)] __attribute__((__section__(".nocache")));

static void i2c_ambiq_callback(void *callback_ctxt, uint32_t status)
{
	const struct device *dev = callback_ctxt;
	struct i2c_ambiq_data *data = dev->data;

	if (data->callback) {
		data->callback(dev, status, data->callback_data);
	}
	data->transfer_status = status;
}
#endif

static void i2c_ambiq_isr(const struct device *dev)
{
	uint32_t ui32Status;
	struct i2c_ambiq_data *data = dev->data;

	am_hal_iom_interrupt_status_get(data->iom_handler, false, &ui32Status);
	am_hal_iom_interrupt_clear(data->iom_handler, ui32Status);
	am_hal_iom_interrupt_service(data->iom_handler, ui32Status);
	k_sem_give(&data->transfer_sem);
}

static int i2c_ambiq_read(const struct device *dev, struct i2c_msg *msg, uint16_t addr)
{
	struct i2c_ambiq_data *data = dev->data;

	int ret = 0;

	am_hal_iom_transfer_t trans = {0};

	trans.ui8Priority = 1;
	trans.eDirection = AM_HAL_IOM_RX;
	trans.uPeerInfo.ui32I2CDevAddr = addr;
	trans.ui32NumBytes = msg->len;
	trans.pui32RxBuffer = (uint32_t *)msg->buf;

#ifdef CONFIG_I2C_AMBIQ_DMA
	data->transfer_status = -EFAULT;
	ret = am_hal_iom_nonblocking_transfer(data->iom_handler, &trans, i2c_ambiq_callback,
					      (void *)dev);
	if (k_sem_take(&data->transfer_sem, K_MSEC(I2C_TRANSFER_TIMEOUT_MSEC))) {
		LOG_ERR("Timeout waiting for transfer complete");
		/* cancel timed out transaction */
		am_hal_iom_disable(data->iom_handler);
		/* clean up for next xfer */
		k_sem_reset(&data->transfer_sem);
		am_hal_iom_enable(data->iom_handler);
		return -ETIMEDOUT;
	}
	ret = data->transfer_status;
#else
	ret = am_hal_iom_blocking_transfer(data->iom_handler, &trans);
#endif
	return (ret != AM_HAL_STATUS_SUCCESS) ? -EIO : 0;
}

static int i2c_ambiq_write(const struct device *dev, struct i2c_msg *msg, uint16_t addr)
{
	struct i2c_ambiq_data *data = dev->data;

	int ret = 0;

	am_hal_iom_transfer_t trans = {0};

	trans.ui8Priority = 1;
	trans.eDirection = AM_HAL_IOM_TX;
	trans.uPeerInfo.ui32I2CDevAddr = addr;
	trans.ui32NumBytes = msg->len;
	trans.pui32TxBuffer = (uint32_t *)msg->buf;

#ifdef CONFIG_I2C_AMBIQ_DMA
	data->transfer_status = -EFAULT;
	ret = am_hal_iom_nonblocking_transfer(data->iom_handler, &trans, i2c_ambiq_callback,
					      (void *)dev);

	if (k_sem_take(&data->transfer_sem, K_MSEC(I2C_TRANSFER_TIMEOUT_MSEC))) {
		LOG_ERR("Timeout waiting for transfer complete");
		/* cancel timed out transaction */
		am_hal_iom_disable(data->iom_handler);
		/* clean up for next xfer */
		k_sem_reset(&data->transfer_sem);
		am_hal_iom_enable(data->iom_handler);
		return -ETIMEDOUT;
	}
	ret = data->transfer_status;
#else
	ret = am_hal_iom_blocking_transfer(data->iom_handler, &trans);
#endif

	return (ret != AM_HAL_STATUS_SUCCESS) ? -EIO : 0;
}

static int i2c_ambiq_configure(const struct device *dev, uint32_t dev_config)
{
	struct i2c_ambiq_data *data = dev->data;

	if (!(I2C_MODE_CONTROLLER & dev_config)) {
		return -EINVAL;
	}

	switch (I2C_SPEED_GET(dev_config)) {
	case I2C_SPEED_STANDARD:
		data->iom_cfg.ui32ClockFreq = AM_HAL_IOM_100KHZ;
		break;
	case I2C_SPEED_FAST:
		data->iom_cfg.ui32ClockFreq = AM_HAL_IOM_400KHZ;
		break;
	case I2C_SPEED_FAST_PLUS:
		data->iom_cfg.ui32ClockFreq = AM_HAL_IOM_1MHZ;
		break;
	default:
		return -EINVAL;
	}

#ifdef CONFIG_I2C_AMBIQ_DMA
	data->iom_cfg.pNBTxnBuf = i2c_dma_tcb_buf[data->inst_idx].buf;
	data->iom_cfg.ui32NBTxnBufLength = CONFIG_I2C_DMA_TCB_BUFFER_SIZE;
#endif

	am_hal_iom_configure(data->iom_handler, &data->iom_cfg);

	return 0;
}

static int i2c_ambiq_transfer(const struct device *dev, struct i2c_msg *msgs, uint8_t num_msgs,
			      uint16_t addr)
{
	struct i2c_ambiq_data *data = dev->data;
	int ret = 0;

	if (!num_msgs) {
		return 0;
	}

	i2c_ambiq_pm_policy_state_lock_get(dev);

	/* Send out messages */
	k_sem_take(&data->bus_sem, K_FOREVER);

	for (int i = 0; i < num_msgs; i++) {
		if (msgs[i].flags & I2C_MSG_READ) {
			ret = i2c_ambiq_read(dev, &(msgs[i]), addr);
		} else {
			ret = i2c_ambiq_write(dev, &(msgs[i]), addr);
		}

		if (ret != 0) {
			LOG_ERR("i2c transfer failed: %d", ret);
			break;
		}
	}

	k_sem_give(&data->bus_sem);

	i2c_ambiq_pm_policy_state_lock_put(dev);

	return ret;
}

#if CONFIG_I2C_AMBIQ_BUS_RECOVERY
static void i2c_ambiq_bitbang_set_scl(void *io_context, int state)
{
	const struct i2c_ambiq_config *config = io_context;

	gpio_pin_set_dt(&config->scl, state);
}

static void i2c_ambiq_bitbang_set_sda(void *io_context, int state)
{
	const struct i2c_ambiq_config *config = io_context;

	gpio_pin_set_dt(&config->sda, state);
}

static int i2c_ambiq_bitbang_get_sda(void *io_context)
{
	const struct i2c_ambiq_config *config = io_context;

	return gpio_pin_get_dt(&config->sda) == 0 ? 0 : 1;
}

static int i2c_ambiq_recover_bus(const struct device *dev)
{
	const struct i2c_ambiq_config *config = dev->config;
	struct i2c_ambiq_data *data = dev->data;
	struct i2c_bitbang bitbang_ctx;
	struct i2c_bitbang_io bitbang_io = {
		.set_scl = i2c_ambiq_bitbang_set_scl,
		.set_sda = i2c_ambiq_bitbang_set_sda,
		.get_sda = i2c_ambiq_bitbang_get_sda,
	};
	uint32_t bitrate_cfg;
	int error = 0;

	LOG_ERR("attempting to recover bus");

	if (!gpio_is_ready_dt(&config->scl)) {
		LOG_ERR("SCL GPIO device not ready");
		return -EIO;
	}

	if (!gpio_is_ready_dt(&config->sda)) {
		LOG_ERR("SDA GPIO device not ready");
		return -EIO;
	}

	k_sem_take(&data->bus_sem, K_FOREVER);

	error = gpio_pin_configure_dt(&config->scl, GPIO_OUTPUT_HIGH);
	if (error != 0) {
		LOG_ERR("failed to configure SCL GPIO (err %d)", error);
		goto restore;
	}

	error = gpio_pin_configure_dt(&config->sda, GPIO_OUTPUT_HIGH);
	if (error != 0) {
		LOG_ERR("failed to configure SDA GPIO (err %d)", error);
		goto restore;
	}

	i2c_bitbang_init(&bitbang_ctx, &bitbang_io, (void *)config);

	bitrate_cfg = i2c_map_dt_bitrate(config->bitrate) | I2C_MODE_CONTROLLER;
	error = i2c_bitbang_configure(&bitbang_ctx, bitrate_cfg);
	if (error != 0) {
		LOG_ERR("failed to configure I2C bitbang (err %d)", error);
		goto restore;
	}

	error = i2c_bitbang_recover_bus(&bitbang_ctx);
	if (error != 0) {
		LOG_ERR("failed to recover bus (err %d)", error);
	}

restore:
	(void)pinctrl_apply_state(config->pcfg, PINCTRL_STATE_DEFAULT);

	k_sem_give(&data->bus_sem);

	return error;
}
#endif /* CONFIG_I2C_AMBIQ_BUS_RECOVERY */

static int i2c_ambiq_init(const struct device *dev)
{
	struct i2c_ambiq_data *data = dev->data;
	const struct i2c_ambiq_config *config = dev->config;
	uint32_t bitrate_cfg = i2c_map_dt_bitrate(config->bitrate);
	int ret = 0;

	data->iom_cfg.eInterfaceMode = AM_HAL_IOM_I2C_MODE;

	if (AM_HAL_STATUS_SUCCESS !=
	    am_hal_iom_initialize((config->base - IOM0_BASE) / config->size,
				  &data->iom_handler)) {
		LOG_ERR("Fail to initialize I2C\n");
		return -ENXIO;
	}

	ret = config->pwr_func();

	ret |= i2c_ambiq_configure(dev, I2C_MODE_CONTROLLER | bitrate_cfg);
	if (ret < 0) {
		LOG_ERR("Fail to config I2C\n");
		goto end;
	}

	ret = pinctrl_apply_state(config->pcfg, PINCTRL_STATE_DEFAULT);
	if (ret < 0) {
		LOG_ERR("Fail to config I2C pins\n");
		goto end;
	}

#ifdef CONFIG_I2C_AMBIQ_DMA
	am_hal_iom_interrupt_clear(data->iom_handler, AM_HAL_IOM_INT_DCMP | AM_HAL_IOM_INT_CMDCMP);
	am_hal_iom_interrupt_enable(data->iom_handler, AM_HAL_IOM_INT_DCMP | AM_HAL_IOM_INT_CMDCMP);
	config->irq_config_func();
#endif

	if (AM_HAL_STATUS_SUCCESS != am_hal_iom_enable(data->iom_handler)) {
		LOG_ERR("Fail to enable I2C\n");
		ret = -EIO;
	}
end:
	if (ret < 0) {
		am_hal_iom_uninitialize(data->iom_handler);
	}
	return ret;
}

static const struct i2c_driver_api i2c_ambiq_driver_api = {
	.configure = i2c_ambiq_configure,
	.transfer = i2c_ambiq_transfer,
#if CONFIG_I2C_AMBIQ_BUS_RECOVERY
	.recover_bus = i2c_ambiq_recover_bus,
#endif /* CONFIG_I2C_AMBIQ_BUS_RECOVERY */
#ifdef CONFIG_I2C_RTIO
	.iodev_submit = i2c_iodev_submit_fallback,
#endif
};

#ifdef CONFIG_PM_DEVICE
static int i2c_ambiq_pm_action(const struct device *dev, enum pm_device_action action)
{
	struct i2c_ambiq_data *data = dev->data;
	uint32_t ret;
	am_hal_sysctrl_power_state_e status;

	switch (action) {
	case PM_DEVICE_ACTION_RESUME:
		status = AM_HAL_SYSCTRL_WAKE;
		break;
	case PM_DEVICE_ACTION_SUSPEND:
		status = AM_HAL_SYSCTRL_DEEPSLEEP;
		break;
	default:
		return -ENOTSUP;
	}

	ret = am_hal_iom_power_ctrl(data->iom_handler, status, true);

	if (ret != AM_HAL_STATUS_SUCCESS) {
		return -EPERM;
	} else {
		return 0;
	}
}
#endif /* CONFIG_PM_DEVICE */

#define AMBIQ_I2C_DEFINE(n)                                                                        \
	PINCTRL_DT_INST_DEFINE(n);                                                                 \
	static int pwr_on_ambiq_i2c_##n(void)                                                      \
	{                                                                                          \
		uint32_t addr = DT_REG_ADDR(DT_INST_PHANDLE(n, ambiq_pwrcfg)) +                    \
				DT_INST_PHA(n, ambiq_pwrcfg, offset);                              \
		sys_write32((sys_read32(addr) | DT_INST_PHA(n, ambiq_pwrcfg, mask)), addr);        \
		k_busy_wait(PWRCTRL_MAX_WAIT_US);                                                  \
		return 0;                                                                          \
	}                                                                                          \
	static void i2c_irq_config_func_##n(void)                                                  \
	{                                                                                          \
		IRQ_CONNECT(DT_INST_IRQN(n), DT_INST_IRQ(n, priority), i2c_ambiq_isr,              \
			    DEVICE_DT_INST_GET(n), 0);                                             \
		irq_enable(DT_INST_IRQN(n));                                                       \
	};                                                                                         \
	static struct i2c_ambiq_data i2c_ambiq_data##n = {                                         \
		.bus_sem = Z_SEM_INITIALIZER(i2c_ambiq_data##n.bus_sem, 1, 1),                     \
		.transfer_sem = Z_SEM_INITIALIZER(i2c_ambiq_data##n.transfer_sem, 0, 1),           \
		.inst_idx = n,                                                                     \
	};                                                                                         \
	static const struct i2c_ambiq_config i2c_ambiq_config##n = {                               \
		.base = DT_INST_REG_ADDR(n),                                                       \
		.size = DT_INST_REG_SIZE(n),                                                       \
		.bitrate = DT_INST_PROP(n, clock_frequency),                                       \
		.pcfg = PINCTRL_DT_INST_DEV_CONFIG_GET(n),                                         \
		.irq_config_func = i2c_irq_config_func_##n,                                        \
		.pwr_func = pwr_on_ambiq_i2c_##n,                                                  \
		IF_ENABLED(CONFIG_I2C_AMBIQ_BUS_RECOVERY,			\
		(.scl = GPIO_DT_SPEC_INST_GET_OR(n, scl_gpios, {0}),\
		 .sda = GPIO_DT_SPEC_INST_GET_OR(n, sda_gpios, {0}),)) };      \
	PM_DEVICE_DT_INST_DEFINE(n, i2c_ambiq_pm_action);                                          \
	I2C_DEVICE_DT_INST_DEFINE(n, i2c_ambiq_init, PM_DEVICE_DT_INST_GET(n), &i2c_ambiq_data##n, \
				  &i2c_ambiq_config##n, POST_KERNEL, CONFIG_I2C_INIT_PRIORITY,     \
				  &i2c_ambiq_driver_api);

DT_INST_FOREACH_STATUS_OKAY(AMBIQ_I2C_DEFINE)