// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Troy Mitchell <troy.mitchell@linux.spacemit.com> */

#include <linux/bitfield.h>
#include <linux/clk.h>
#include <linux/reset.h>
#include <sound/dmaengine_pcm.h>
#include <sound/pcm.h>
#include <sound/pcm_params.h>

#define SSCR                    0x00    /* SPI/I2S top control register */
#define SSFCR                   0x04    /* SPI/I2S FIFO control register */
#define SSINTEN                 0x08    /* SPI/I2S interrupt enable register */
#define SSDATR                  0x10    /* SPI/I2S data register */
#define SSPSP                   0x18    /* SPI/I2S programmable serial protocol control register */
#define SSRWT                   0x24    /* SPI/I2S root control register */

/* SPI/I2S Work data size, register bits value 0~31 indicated data size 1~32 bits */
#define SSCR_FIELD_DSS          GENMASK(9, 5)
#define SSCR_DW_8BYTE           FIELD_PREP(SSCR_FIELD_DSS, 0x7)
#define SSCR_DW_16BYTE          FIELD_PREP(SSCR_FIELD_DSS, 0xf)
#define SSCR_DW_18BYTE          FIELD_PREP(SSCR_FIELD_DSS, 0x11)
#define SSCR_DW_32BYTE          FIELD_PREP(SSCR_FIELD_DSS, 0x1f)

#define SSCR_SSE                BIT(0)          /* SPI/I2S Enable */
#define SSCR_FRF_PSP            GENMASK(2, 1)   /* Frame Format*/
#define SSCR_TRAIL              BIT(13)         /* Trailing Byte */

#define SSFCR_FIELD_TFT         GENMASK(3, 0)   /* TXFIFO Trigger Threshold */
#define SSFCR_FIELD_RFT         GENMASK(8, 5)   /* RXFIFO Trigger Threshold */
#define SSFCR_TSRE              BIT(10)         /* Transmit Service Request Enable */
#define SSFCR_RSRE              BIT(11)         /* Receive Service Request Enable */

#define SSPSP_FSRT              BIT(3)          /* Frame Sync Relative Timing Bit */
#define SSPSP_SFRMP             BIT(4)          /* Serial Frame Polarity */
#define SSPSP_FIELD_SFRMWDTH    GENMASK(17, 12) /* Serial Frame Width field  */

#define SSRWT_RWOT              BIT(0)          /* Receive Without Transmit */

#define SPACEMIT_PCM_RATES      (SNDRV_PCM_RATE_8000 | SNDRV_PCM_RATE_16000 | \
                                SNDRV_PCM_RATE_48000)
#define SPACEMIT_PCM_FORMATS    (SNDRV_PCM_FMTBIT_S16_LE | SNDRV_PCM_FMTBIT_S32_LE)

#define SPACEMIT_I2S_PERIOD_SIZE 1024

struct spacemit_i2s_dev {
        struct device *dev;

        void __iomem *base;

        struct reset_control *reset;

        struct clk *sysclk;
        struct clk *bclk;
        struct clk *sspa_clk;

        struct snd_dmaengine_dai_dma_data capture_dma_data;
        struct snd_dmaengine_dai_dma_data playback_dma_data;

        bool has_capture;
        bool has_playback;

        int dai_fmt;

        int started_count;
};

static const struct snd_pcm_hardware spacemit_pcm_hardware = {
        .info             = SNDRV_PCM_INFO_INTERLEAVED |
                            SNDRV_PCM_INFO_BATCH,
        .formats          = SPACEMIT_PCM_FORMATS,
        .rates            = SPACEMIT_PCM_RATES,
        .rate_min         = SNDRV_PCM_RATE_8000,
        .rate_max         = SNDRV_PCM_RATE_192000,
        .channels_min     = 1,
        .channels_max     = 2,
        .buffer_bytes_max = SPACEMIT_I2S_PERIOD_SIZE * 4 * 4,
        .period_bytes_min = SPACEMIT_I2S_PERIOD_SIZE * 2,
        .period_bytes_max = SPACEMIT_I2S_PERIOD_SIZE * 4,
        .periods_min      = 2,
        .periods_max      = 4,
};

static const struct snd_dmaengine_pcm_config spacemit_dmaengine_pcm_config = {
        .pcm_hardware = &spacemit_pcm_hardware,
        .prepare_slave_config = snd_dmaengine_pcm_prepare_slave_config,
        .chan_names = {"tx", "rx"},
        .prealloc_buffer_size = 32 * 1024,
};

static void spacemit_i2s_init(struct spacemit_i2s_dev *i2s)
{
        u32 sscr_val, sspsp_val, ssfcr_val, ssrwt_val;

        sscr_val = SSCR_TRAIL | SSCR_FRF_PSP;
        ssfcr_val = FIELD_PREP(SSFCR_FIELD_TFT, 5) |
                    FIELD_PREP(SSFCR_FIELD_RFT, 5) |
                    SSFCR_RSRE | SSFCR_TSRE;
        ssrwt_val = SSRWT_RWOT;
        sspsp_val = SSPSP_SFRMP;

        writel(sscr_val, i2s->base + SSCR);
        writel(ssfcr_val, i2s->base + SSFCR);
        writel(sspsp_val, i2s->base + SSPSP);
        writel(ssrwt_val, i2s->base + SSRWT);
        writel(0, i2s->base + SSINTEN);
}

static int spacemit_i2s_hw_params(struct snd_pcm_substream *substream,
                                  struct snd_pcm_hw_params *params,
                                  struct snd_soc_dai *dai)
{
        struct spacemit_i2s_dev *i2s = snd_soc_dai_get_drvdata(dai);
        struct snd_dmaengine_dai_dma_data *dma_data;
        u32 data_width, data_bits;
        unsigned long bclk_rate;
        u32 val;
        int ret;

        val = readl(i2s->base + SSCR);
        if (val & SSCR_SSE)
                return 0;

        dma_data = &i2s->playback_dma_data;

        if (substream->stream == SNDRV_PCM_STREAM_CAPTURE)
                dma_data = &i2s->capture_dma_data;

        switch (params_format(params)) {
        case SNDRV_PCM_FORMAT_S8:
                data_bits = 8;
                data_width = SSCR_DW_8BYTE;
                dma_data->maxburst = 8;
                dma_data->addr_width = DMA_SLAVE_BUSWIDTH_1_BYTE;
                break;
        case SNDRV_PCM_FORMAT_S16_LE:
                data_bits = 16;
                data_width = SSCR_DW_16BYTE;
                dma_data->maxburst = 16;
                dma_data->addr_width = DMA_SLAVE_BUSWIDTH_2_BYTES;
                break;
        case SNDRV_PCM_FORMAT_S32_LE:
                data_bits = 32;
                data_width = SSCR_DW_32BYTE;
                dma_data->maxburst = 32;
                dma_data->addr_width = DMA_SLAVE_BUSWIDTH_4_BYTES;
                break;
        default:
                dev_dbg(i2s->dev, "unexpected data width type");
                return -EINVAL;
        }

        switch (i2s->dai_fmt & SND_SOC_DAIFMT_FORMAT_MASK) {
        case SND_SOC_DAIFMT_I2S:
                if (data_bits == 16) {
                        data_width = SSCR_DW_32BYTE;
                        dma_data->maxburst = 32;
                        dma_data->addr_width = DMA_SLAVE_BUSWIDTH_4_BYTES;
                }

                snd_pcm_hw_constraint_minmax(substream->runtime,
                                             SNDRV_PCM_HW_PARAM_CHANNELS,
                                             1, 2);
                snd_pcm_hw_constraint_mask64(substream->runtime,
                                             SNDRV_PCM_HW_PARAM_FORMAT,
                                             SNDRV_PCM_FMTBIT_S16_LE);
                break;
        case SND_SOC_DAIFMT_DSP_A:
        case SND_SOC_DAIFMT_DSP_B:
                snd_pcm_hw_constraint_minmax(substream->runtime,
                                             SNDRV_PCM_HW_PARAM_CHANNELS,
                                             1, 1);
                snd_pcm_hw_constraint_mask64(substream->runtime,
                                             SNDRV_PCM_HW_PARAM_FORMAT,
                                             SNDRV_PCM_FMTBIT_S32_LE);
                break;
        default:
                dev_dbg(i2s->dev, "unexpected format type");
                return -EINVAL;

        }

        val = readl(i2s->base + SSCR);
        val &= ~SSCR_DW_32BYTE;
        val |= data_width;
        writel(val, i2s->base + SSCR);

        bclk_rate = params_channels(params) *
                    params_rate(params) *
                    data_bits;

        ret = clk_set_rate(i2s->bclk, bclk_rate);
        if (ret)
                return ret;

        return clk_set_rate(i2s->sspa_clk, bclk_rate);
}

static int spacemit_i2s_set_sysclk(struct snd_soc_dai *cpu_dai, int clk_id,
                                   unsigned int freq, int dir)
{
        struct spacemit_i2s_dev *i2s = dev_get_drvdata(cpu_dai->dev);

        if (freq == 0)
                return 0;

        return clk_set_rate(i2s->sysclk, freq);
}

static int spacemit_i2s_set_fmt(struct snd_soc_dai *cpu_dai,
                                unsigned int fmt)
{
        struct spacemit_i2s_dev *i2s = dev_get_drvdata(cpu_dai->dev);
        u32 sspsp_val;

        sspsp_val = readl(i2s->base + SSPSP);
        sspsp_val &= ~SSPSP_FIELD_SFRMWDTH;
        sspsp_val |= SSPSP_FSRT;

        i2s->dai_fmt = fmt;

        switch (fmt & SND_SOC_DAIFMT_FORMAT_MASK) {
        case SND_SOC_DAIFMT_I2S:
                sspsp_val |= FIELD_PREP(SSPSP_FIELD_SFRMWDTH, 0x10);
                break;
        case SND_SOC_DAIFMT_DSP_B:
                /* DSP_B: next frame asserted after previous frame end, so clear FSRT */
                sspsp_val &= ~SSPSP_FSRT;
                fallthrough;
        case SND_SOC_DAIFMT_DSP_A:
                sspsp_val |= FIELD_PREP(SSPSP_FIELD_SFRMWDTH, 0x1);
                break;
        default:
                dev_dbg(i2s->dev, "unexpected format type");
                return -EINVAL;
        }

        writel(sspsp_val, i2s->base + SSPSP);

        return 0;
}

static int spacemit_i2s_trigger(struct snd_pcm_substream *substream,
                                int cmd, struct snd_soc_dai *dai)
{
        struct spacemit_i2s_dev *i2s = snd_soc_dai_get_drvdata(dai);
        u32 val;

        switch (cmd) {
        case SNDRV_PCM_TRIGGER_START:
        case SNDRV_PCM_TRIGGER_RESUME:
        case SNDRV_PCM_TRIGGER_PAUSE_RELEASE:
                if (!i2s->started_count) {
                        val = readl(i2s->base + SSCR);
                        val |= SSCR_SSE;
                        writel(val, i2s->base + SSCR);
                }
                i2s->started_count++;
                break;
        case SNDRV_PCM_TRIGGER_STOP:
        case SNDRV_PCM_TRIGGER_SUSPEND:
        case SNDRV_PCM_TRIGGER_PAUSE_PUSH:
                if (i2s->started_count)
                        i2s->started_count--;

                if (!i2s->started_count) {
                        val = readl(i2s->base + SSCR);
                        val &= ~SSCR_SSE;
                        writel(val, i2s->base + SSCR);
                }
                break;
        default:
                return -EINVAL;
        }

        return 0;
}

static int spacemit_i2s_dai_probe(struct snd_soc_dai *dai)
{
        struct spacemit_i2s_dev *i2s = snd_soc_dai_get_drvdata(dai);

        snd_soc_dai_init_dma_data(dai,
                                  i2s->has_playback ? &i2s->playback_dma_data : NULL,
                                  i2s->has_capture ? &i2s->capture_dma_data : NULL);

        reset_control_deassert(i2s->reset);

        spacemit_i2s_init(i2s);

        return 0;
}

static int spacemit_i2s_dai_remove(struct snd_soc_dai *dai)
{
        struct spacemit_i2s_dev *i2s = snd_soc_dai_get_drvdata(dai);

        reset_control_assert(i2s->reset);

        return 0;
}

static const struct snd_soc_dai_ops spacemit_i2s_dai_ops = {
        .probe = spacemit_i2s_dai_probe,
        .remove = spacemit_i2s_dai_remove,
        .hw_params = spacemit_i2s_hw_params,
        .set_sysclk = spacemit_i2s_set_sysclk,
        .set_fmt = spacemit_i2s_set_fmt,
        .trigger = spacemit_i2s_trigger,
};

static struct snd_soc_dai_driver spacemit_i2s_dai = {
        .ops = &spacemit_i2s_dai_ops,
        .playback = {
                .channels_min = 1,
                .channels_max = 2,
                .rates = SPACEMIT_PCM_RATES,
                .rate_min = SNDRV_PCM_RATE_8000,
                .rate_max = SNDRV_PCM_RATE_48000,
                .formats = SPACEMIT_PCM_FORMATS,
        },
        .capture = {
                .channels_min = 1,
                .channels_max = 2,
                .rates = SPACEMIT_PCM_RATES,
                .rate_min = SNDRV_PCM_RATE_8000,
                .rate_max = SNDRV_PCM_RATE_48000,
                .formats = SPACEMIT_PCM_FORMATS,
        },
        .symmetric_rate = 1,
};

static int spacemit_i2s_init_dai(struct spacemit_i2s_dev *i2s,
                                 struct snd_soc_dai_driver **dp,
                                 dma_addr_t addr)
{
        struct device_node *node = i2s->dev->of_node;
        struct snd_soc_dai_driver *dai;
        struct property *dma_names;
        const char *dma_name;

        of_property_for_each_string(node, "dma-names", dma_names, dma_name) {
                if (!strcmp(dma_name, "tx"))
                        i2s->has_playback = true;
                if (!strcmp(dma_name, "rx"))
                        i2s->has_capture = true;
        }

        dai = devm_kmemdup(i2s->dev, &spacemit_i2s_dai,
                           sizeof(*dai), GFP_KERNEL);
        if (!dai)
                return -ENOMEM;

        if (i2s->has_playback) {
                dai->playback.stream_name = "Playback";
                dai->playback.channels_min = 1;
                dai->playback.channels_max = 2;
                dai->playback.rates = SPACEMIT_PCM_RATES;
                dai->playback.formats = SPACEMIT_PCM_FORMATS;

                i2s->playback_dma_data.addr_width = DMA_SLAVE_BUSWIDTH_2_BYTES;
                i2s->playback_dma_data.maxburst = 32;
                i2s->playback_dma_data.addr = addr;
        }

        if (i2s->has_capture) {
                dai->capture.stream_name = "Capture";
                dai->capture.channels_min = 1;
                dai->capture.channels_max = 2;
                dai->capture.rates = SPACEMIT_PCM_RATES;
                dai->capture.formats = SPACEMIT_PCM_FORMATS;

                i2s->capture_dma_data.addr_width = DMA_SLAVE_BUSWIDTH_2_BYTES;
                i2s->capture_dma_data.maxburst = 32;
                i2s->capture_dma_data.addr = addr;
        }

        if (dp)
                *dp = dai;

        return 0;
}

static const struct snd_soc_component_driver spacemit_i2s_component = {
        .name = "i2s-k1",
        .legacy_dai_naming = 1,
};

static int spacemit_i2s_probe(struct platform_device *pdev)
{
        struct snd_soc_dai_driver *dai;
        struct spacemit_i2s_dev *i2s;
        struct resource *res;
        struct clk *clk;
        int ret;

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

        i2s->dev = &pdev->dev;

        i2s->sysclk = devm_clk_get_enabled(i2s->dev, "sysclk");
        if (IS_ERR(i2s->sysclk))
                return dev_err_probe(i2s->dev, PTR_ERR(i2s->sysclk),
                                     "failed to enable sysbase clock\n");

        i2s->bclk = devm_clk_get_enabled(i2s->dev, "bclk");
        if (IS_ERR(i2s->bclk))
                return dev_err_probe(i2s->dev, PTR_ERR(i2s->bclk), "failed to enable bit clock\n");

        clk = devm_clk_get_enabled(i2s->dev, "sspa_bus");
        if (IS_ERR(clk))
                return dev_err_probe(i2s->dev, PTR_ERR(clk), "failed to enable sspa_bus clock\n");

        i2s->sspa_clk = devm_clk_get_enabled(i2s->dev, "sspa");
        if (IS_ERR(i2s->sspa_clk))
                return dev_err_probe(i2s->dev, PTR_ERR(i2s->sspa_clk),
                                     "failed to enable sspa clock\n");

        i2s->base = devm_platform_get_and_ioremap_resource(pdev, 0, &res);
        if (IS_ERR(i2s->base))
                return dev_err_probe(i2s->dev, PTR_ERR(i2s->base), "failed to map registers\n");

        i2s->reset = devm_reset_control_get_exclusive(&pdev->dev, NULL);
        if (IS_ERR(i2s->reset))
                return dev_err_probe(i2s->dev, PTR_ERR(i2s->reset),
                                     "failed to get reset control");

        dev_set_drvdata(i2s->dev, i2s);

        ret = spacemit_i2s_init_dai(i2s, &dai, res->start + SSDATR);
        if (ret)
                return ret;

        ret = devm_snd_soc_register_component(i2s->dev,
                                              &spacemit_i2s_component,
                                              dai, 1);
        if (ret)
                return dev_err_probe(i2s->dev, ret, "failed to register component");

        return devm_snd_dmaengine_pcm_register(&pdev->dev, &spacemit_dmaengine_pcm_config, 0);
}

static const struct of_device_id spacemit_i2s_of_match[] = {
        { .compatible = "spacemit,k1-i2s", },
        { /* sentinel */ }
};
MODULE_DEVICE_TABLE(of, spacemit_i2s_of_match);

static struct platform_driver spacemit_i2s_driver = {
        .probe = spacemit_i2s_probe,
        .driver = {
                .name = "i2s-k1",
                .of_match_table = spacemit_i2s_of_match,
        },
};
module_platform_driver(spacemit_i2s_driver);

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("I2S bus driver for SpacemiT K1 SoC");