---

#define DT_DRV_COMPAT nxp_mcux_i2s

#include <errno.h>
#include <string.h>
#include <zephyr/sys/__assert.h>
#include <zephyr/kernel.h>
#include <zephyr/device.h>
#include <zephyr/init.h>
#include <zephyr/drivers/dma.h>
#include <zephyr/drivers/i2s.h>
#include <zephyr/drivers/pinctrl.h>
#include <zephyr/drivers/clock_control.h>
#include <zephyr/dt-bindings/clock/imx_ccm.h>
#include <zephyr/sys/barrier.h>
#include <zephyr/device.h>
#include <soc.h>

#include <fsl_sai.h>
#include <fsl_edma.h>

#include <zephyr/logging/log.h>
LOG_MODULE_REGISTER(dev_i2s_mcux, CONFIG_I2S_LOG_LEVEL);

#define NUM_DMA_BLOCKS_RX_PREP 3
#if defined(CONFIG_DMA_MCUX_EDMA)
BUILD_ASSERT(NUM_DMA_BLOCKS_RX_PREP >= 3,
	     "eDMA avoids TCD coherency issue if NUM_DMA_BLOCKS_RX_PREP >= 3");
#endif /* CONFIG_DMA_MCUX_EDMA */

#define MAX_TX_DMA_BLOCKS CONFIG_DMA_TCD_QUEUE_SIZE
BUILD_ASSERT(MAX_TX_DMA_BLOCKS > NUM_DMA_BLOCKS_RX_PREP,
	     "NUM_DMA_BLOCKS_RX_PREP must be < CONFIG_DMA_TCD_QUEUE_SIZE");

#define SAI_WORD_SIZE_BITS_MIN 8
#define SAI_WORD_SIZE_BITS_MAX 32

#define SAI_WORD_PER_FRAME_MIN 0
#define SAI_WORD_PER_FRAME_MAX 32

/*
 * SAI driver uses source_gather_en/dest_scatter_en feature of DMA, and relies
 * on DMA driver managing circular list of DMA blocks.  Like eDMA driver links
 * Transfer Control Descriptors (TCDs) in list, and manages the tcdpool.
 * Calling dma_reload() adds new DMA block to DMA channel already configured,
 * into the DMA driver's circular list of blocks.

 * This indicates the Tx/Rx stream.
 *
 * in_queue and out_queue are used as follows
 *   transmit stream:
 *   application provided buffer is queued to in_queue until loaded to DMA.
 *   when DMA channel is idle, buffer is retrieved from in_queue and loaded
 *   to DMA and queued to out_queue. when DMA completes, buffer is retrieved
 *   from out_queue and freed.
 *
 *   receive stream:
 *   driver allocates buffer from slab and loads DMA buffer is queued to
 *   in_queue when DMA completes, buffer is retrieved from in_queue
 *   and queued to out_queue when application reads, buffer is read
 *   (may optionally block) from out_queue and presented to application.
 */
struct stream {
	enum i2s_state state;
	uint32_t dma_channel;
	uint32_t start_channel;
	void (*irq_call_back)(void);
	struct i2s_config cfg;
	struct dma_config dma_cfg;
	struct dma_block_config dma_block;
	uint8_t free_tx_dma_blocks;
	bool last_block;
	struct k_msgq in_queue;
	struct k_msgq out_queue;
};

struct i2s_mcux_config {
	I2S_Type *base;
	uint32_t clk_src;
	uint32_t clk_pre_div;
	uint32_t clk_src_div;
	uint32_t pll_src;
	uint32_t pll_lp;
	uint32_t pll_pd;
	uint32_t pll_num;
	uint32_t pll_den;
	uint32_t mclk_control_base;
	uint32_t mclk_pin_mask;
	uint32_t mclk_pin_offset;
	bool mclk_output;
	uint32_t tx_channel;
	clock_control_subsys_t clk_sub_sys;
	const struct device *ccm_dev;
	const struct pinctrl_dev_config *pinctrl;
	void (*irq_connect)(const struct device *dev);
	sai_sync_mode_t rx_sync_mode;
	sai_sync_mode_t tx_sync_mode;
};

/* Device run time data */
struct i2s_dev_data {
	const struct device *dev_dma;
	struct stream tx;
	void *tx_in_msgs[CONFIG_I2S_TX_BLOCK_COUNT];
	void *tx_out_msgs[CONFIG_I2S_TX_BLOCK_COUNT];
	struct stream rx;
	void *rx_in_msgs[CONFIG_I2S_RX_BLOCK_COUNT];
	void *rx_out_msgs[CONFIG_I2S_RX_BLOCK_COUNT];
};

static void i2s_purge_stream_buffers(struct stream *strm, struct k_mem_slab *mem_slab, bool in_drop,
				     bool out_drop)
{
	void *buffer;

	if (in_drop) {
		while (k_msgq_get(&strm->in_queue, &buffer, K_NO_WAIT) == 0) {
			k_mem_slab_free(mem_slab, buffer);
		}
	}

	if (out_drop) {
		while (k_msgq_get(&strm->out_queue, &buffer, K_NO_WAIT) == 0) {
			k_mem_slab_free(mem_slab, buffer);
		}
	}
}

static void i2s_tx_stream_disable(const struct device *dev, bool drop)
{
	struct i2s_dev_data *dev_data = dev->data;
	struct stream *strm = &dev_data->tx;
	const struct device *dev_dma = dev_data->dev_dma;
	const struct i2s_mcux_config *dev_cfg = dev->config;

	LOG_DBG("Stopping DMA channel %u for TX stream", strm->dma_channel);

	/* Disable FIFO DMA request */
	SAI_TxEnableDMA(dev_cfg->base, kSAI_FIFORequestDMAEnable, false);

	dma_stop(dev_dma, strm->dma_channel);

	/* wait for TX FIFO to drain before disabling */
	while ((dev_cfg->base->TCSR & I2S_TCSR_FWF_MASK) == 0) {
		;
	}

	/* Disable the channel FIFO */
	SAI_TxSetChannelFIFOMask(dev_cfg->base, 0);

	/* Disable Tx */
	SAI_TxEnable(dev_cfg->base, false);

	/* If Tx is disabled, reset the FIFO pointer, clear error flags */
	if ((dev_cfg->base->TCSR & I2S_TCSR_TE_MASK) == 0UL) {
		dev_cfg->base->TCSR |= (I2S_TCSR_FR_MASK | I2S_TCSR_SR_MASK);
		dev_cfg->base->TCSR &= ~I2S_TCSR_SR_MASK;
	}

	/* purge buffers queued in the stream */
	i2s_purge_stream_buffers(strm, dev_data->tx.cfg.mem_slab, drop, drop);
}

static void i2s_rx_stream_disable(const struct device *dev, bool in_drop, bool out_drop)
{
	struct i2s_dev_data *dev_data = dev->data;
	struct stream *strm = &dev_data->rx;
	const struct device *dev_dma = dev_data->dev_dma;
	const struct i2s_mcux_config *dev_cfg = dev->config;

	LOG_DBG("Stopping RX stream & DMA channel %u", strm->dma_channel);
	dma_stop(dev_dma, strm->dma_channel);

	/* Disable the channel FIFO */
	SAI_RxSetChannelFIFOMask(dev_cfg->base, 0);

	/* Disable DMA enable bit */
	SAI_RxEnableDMA(dev_cfg->base, kSAI_FIFORequestDMAEnable, false);

	/* Disable Rx */
	SAI_RxEnable(dev_cfg->base, false);

	/* wait for Receiver to disable */
	while (dev_cfg->base->RCSR & I2S_RCSR_RE_MASK) {
		;
	}
	/* reset the FIFO pointer and clear error flags */
	dev_cfg->base->RCSR |= (I2S_RCSR_FR_MASK | I2S_RCSR_SR_MASK);
	dev_cfg->base->RCSR &= ~I2S_RCSR_SR_MASK;

	/* purge buffers queued in the stream */
	i2s_purge_stream_buffers(strm, dev_data->rx.cfg.mem_slab, in_drop, out_drop);
}

static int i2s_tx_reload_multiple_dma_blocks(const struct device *dev, uint8_t *blocks_queued)
{
	struct i2s_dev_data *dev_data = dev->data;
	const struct i2s_mcux_config *dev_cfg = dev->config;
	I2S_Type *base = (I2S_Type *)dev_cfg->base;
	struct stream *strm = &dev_data->tx;
	void *buffer = NULL;
	int ret = 0;
	unsigned int key;

	*blocks_queued = 0;

	key = irq_lock();

	/* queue additional blocks to DMA if in_queue and DMA has free blocks */
	while (strm->free_tx_dma_blocks) {
		/* get the next buffer from queue */
		ret = k_msgq_get(&strm->in_queue, &buffer, K_NO_WAIT);
		if (ret) {
			/* in_queue is empty, no more blocks to send to DMA */
			ret = 0;
			break;
		}

		/* reload the DMA */
		ret = dma_reload(dev_data->dev_dma, strm->dma_channel, (uint32_t)buffer,
				 (uint32_t)&base->TDR[strm->start_channel], strm->cfg.block_size);
		if (ret != 0) {
			LOG_ERR("dma_reload() failed with error 0x%x", ret);
			break;
		}

		(strm->free_tx_dma_blocks)--;

		ret = k_msgq_put(&strm->out_queue, &buffer, K_NO_WAIT);
		if (ret != 0) {
			LOG_ERR("buffer %p -> out %p err %d", buffer, &strm->out_queue, ret);
			break;
		}

		(*blocks_queued)++;
	}

	irq_unlock(key);
	return ret;
}

/* This function is executed in the interrupt context */
static void i2s_dma_tx_callback(const struct device *dma_dev, void *arg, uint32_t channel,
				int status)
{
	const struct device *dev = (struct device *)arg;
	struct i2s_dev_data *dev_data = dev->data;
	struct stream *strm = &dev_data->tx;
	uint8_t blocks_queued;
	void *buffer = NULL;
	int ret;

	LOG_DBG("tx cb");

	ret = k_msgq_get(&strm->out_queue, &buffer, K_NO_WAIT);
	if (ret == 0) {
		/* transmission complete. free the buffer */
		k_mem_slab_free(strm->cfg.mem_slab, buffer);
		(strm->free_tx_dma_blocks)++;
	} else {
		LOG_ERR("no buf in out_queue for channel %u", channel);
	}

	if (strm->free_tx_dma_blocks > MAX_TX_DMA_BLOCKS) {
		strm->state = I2S_STATE_ERROR;
		LOG_ERR("free_tx_dma_blocks exceeded maximum, now %d", strm->free_tx_dma_blocks);
		goto disabled_exit_no_drop;
	}

	/* Received a STOP trigger, terminate TX immediately */
	if (strm->last_block) {
		strm->state = I2S_STATE_READY;
		LOG_DBG("TX STOPPED last_block set");
		goto disabled_exit_no_drop;
	}

	if (ret) {
		/* k_msgq_get() returned error, and was not last_block */
		strm->state = I2S_STATE_ERROR;
		goto disabled_exit_no_drop;
	}

	if (strm->state != I2S_STATE_RUNNING && strm->state != I2S_STATE_STOPPING) {
		goto disabled_exit_drop;
	}

	ret = i2s_tx_reload_multiple_dma_blocks(dev, &blocks_queued);
	if (ret) {
		strm->state = I2S_STATE_ERROR;
		goto disabled_exit_no_drop;
	}

	if (blocks_queued || (strm->free_tx_dma_blocks < MAX_TX_DMA_BLOCKS)) {
		goto enabled_exit;
	}

	/* all DMA blocks are free but no blocks were queued */
	if (strm->state == I2S_STATE_STOPPING) {
		/* TX queue has drained */
		strm->state = I2S_STATE_READY;
		LOG_DBG("TX stream has stopped");
		goto disabled_exit_no_drop;
	}

	LOG_WRN("TX input queue empty!");
	if (strm->free_tx_dma_blocks >= MAX_TX_DMA_BLOCKS) {
		/* In running state, no TX blocks for transferring, so stop
		 * TX (This will disable bit clock to avoid dummy bits
		 * received in RX side.
		 */
		const struct i2s_mcux_config *dev_cfg = dev->config;
		I2S_Type *base = (I2S_Type *)dev_cfg->base;

		SAI_TxEnable(base, false);
		LOG_WRN("TX is paused.");
	}
	goto enabled_exit;


disabled_exit_no_drop:
	i2s_tx_stream_disable(dev, false);
	return;

disabled_exit_drop:
	i2s_tx_stream_disable(dev, true);
	return;

enabled_exit:
	return;
}