/*
 * Copyright 2024 NXP
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */
#include <errno.h>

#include <s32cc-clk-regs.h>

#include <common/debug.h>
#include <drivers/clk.h>
#include <lib/mmio.h>
#include <s32cc-clk-ids.h>
#include <s32cc-clk-modules.h>
#include <s32cc-clk-utils.h>

#define MAX_STACK_DEPTH		(15U)

/* This is used for floating-point precision calculations. */
#define FP_PRECISION		(100000000UL)

struct s32cc_clk_drv {
	uintptr_t fxosc_base;
	uintptr_t armpll_base;
};

static int update_stack_depth(unsigned int *depth)
{
	if (*depth == 0U) {
		return -ENOMEM;
	}

	(*depth)--;
	return 0;
}

static struct s32cc_clk_drv *get_drv(void)
{
	static struct s32cc_clk_drv driver = {
		.fxosc_base = FXOSC_BASE_ADDR,
		.armpll_base = ARMPLL_BASE_ADDR,
	};

	return &driver;
}

static int enable_module(const struct s32cc_clk_obj *module, unsigned int *depth);

static int enable_clk_module(const struct s32cc_clk_obj *module,
			     const struct s32cc_clk_drv *drv,
			     unsigned int *depth)
{
	const struct s32cc_clk *clk = s32cc_obj2clk(module);
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	if (clk == NULL) {
		return -EINVAL;
	}

	if (clk->module != NULL) {
		return enable_module(clk->module, depth);
	}

	if (clk->pclock != NULL) {
		return enable_clk_module(&clk->pclock->desc, drv, depth);
	}

	return -EINVAL;
}

static int get_base_addr(enum s32cc_clk_source id, const struct s32cc_clk_drv *drv,
			 uintptr_t *base)
{
	int ret = 0;

	switch (id) {
	case S32CC_FXOSC:
		*base = drv->fxosc_base;
		break;
	case S32CC_ARM_PLL:
		*base = drv->armpll_base;
		break;
	case S32CC_CGM1:
		ret = -ENOTSUP;
		break;
	case S32CC_FIRC:
		break;
	case S32CC_SIRC:
		break;
	default:
		ret = -EINVAL;
		break;
	}

	if (ret != 0) {
		ERROR("Unknown clock source id: %u\n", id);
	}

	return ret;
}

static void enable_fxosc(const struct s32cc_clk_drv *drv)
{
	uintptr_t fxosc_base = drv->fxosc_base;
	uint32_t ctrl;

	ctrl = mmio_read_32(FXOSC_CTRL(fxosc_base));
	if ((ctrl & FXOSC_CTRL_OSCON) != U(0)) {
		return;
	}

	ctrl = FXOSC_CTRL_COMP_EN;
	ctrl &= ~FXOSC_CTRL_OSC_BYP;
	ctrl |= FXOSC_CTRL_EOCV(0x1);
	ctrl |= FXOSC_CTRL_GM_SEL(0x7);
	mmio_write_32(FXOSC_CTRL(fxosc_base), ctrl);

	/* Switch ON the crystal oscillator. */
	mmio_setbits_32(FXOSC_CTRL(fxosc_base), FXOSC_CTRL_OSCON);

	/* Wait until the clock is stable. */
	while ((mmio_read_32(FXOSC_STAT(fxosc_base)) & FXOSC_STAT_OSC_STAT) == U(0)) {
	}
}

static int enable_osc(const struct s32cc_clk_obj *module,
		      const struct s32cc_clk_drv *drv,
		      unsigned int *depth)
{
	const struct s32cc_osc *osc = s32cc_obj2osc(module);
	int ret = 0;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	switch (osc->source) {
	case S32CC_FXOSC:
		enable_fxosc(drv);
		break;
	/* FIRC and SIRC oscillators are enabled by default */
	case S32CC_FIRC:
		break;
	case S32CC_SIRC:
		break;
	default:
		ERROR("Invalid oscillator %d\n", osc->source);
		ret = -EINVAL;
		break;
	};

	return ret;
}

static int get_pll_mfi_mfn(unsigned long pll_vco, unsigned long ref_freq,
			   uint32_t *mfi, uint32_t *mfn)

{
	unsigned long vco;
	unsigned long mfn64;

	/* FRAC-N mode */
	*mfi = (uint32_t)(pll_vco / ref_freq);

	/* MFN formula : (double)(pll_vco % ref_freq) / ref_freq * 18432.0 */
	mfn64 = pll_vco % ref_freq;
	mfn64 *= FP_PRECISION;
	mfn64 /= ref_freq;
	mfn64 *= 18432UL;
	mfn64 /= FP_PRECISION;

	if (mfn64 > UINT32_MAX) {
		return -EINVAL;
	}

	*mfn = (uint32_t)mfn64;

	vco = ((unsigned long)*mfn * FP_PRECISION) / 18432UL;
	vco += (unsigned long)*mfi * FP_PRECISION;
	vco *= ref_freq;
	vco /= FP_PRECISION;

	if (vco != pll_vco) {
		ERROR("Failed to find MFI and MFN settings for PLL freq %lu. Nearest freq = %lu\n",
		      pll_vco, vco);
		return -EINVAL;
	}

	return 0;
}

static struct s32cc_clkmux *get_pll_mux(const struct s32cc_pll *pll)
{
	const struct s32cc_clk_obj *source = pll->source;
	const struct s32cc_clk *clk;

	if (source == NULL) {
		ERROR("Failed to identify PLL's parent\n");
		return NULL;
	}

	if (source->type != s32cc_clk_t) {
		ERROR("The parent of the PLL isn't a clock\n");
		return NULL;
	}

	clk = s32cc_obj2clk(source);

	if (clk->module == NULL) {
		ERROR("The clock isn't connected to a module\n");
		return NULL;
	}

	source = clk->module;

	if ((source->type != s32cc_clkmux_t) &&
	    (source->type != s32cc_shared_clkmux_t)) {
		ERROR("The parent of the PLL isn't a MUX\n");
		return NULL;
	}

	return s32cc_obj2clkmux(source);
}

static void disable_odiv(uintptr_t pll_addr, uint32_t div_index)
{
	mmio_clrbits_32(PLLDIG_PLLODIV(pll_addr, div_index), PLLDIG_PLLODIV_DE);
}

static void enable_odiv(uintptr_t pll_addr, uint32_t div_index)
{
	mmio_setbits_32(PLLDIG_PLLODIV(pll_addr, div_index), PLLDIG_PLLODIV_DE);
}

static void disable_odivs(uintptr_t pll_addr, uint32_t ndivs)
{
	uint32_t i;

	for (i = 0; i < ndivs; i++) {
		disable_odiv(pll_addr, i);
	}
}

static void enable_pll_hw(uintptr_t pll_addr)
{
	/* Enable the PLL. */
	mmio_write_32(PLLDIG_PLLCR(pll_addr), 0x0);

	/* Poll until PLL acquires lock. */
	while ((mmio_read_32(PLLDIG_PLLSR(pll_addr)) & PLLDIG_PLLSR_LOCK) == 0U) {
	}
}

static void disable_pll_hw(uintptr_t pll_addr)
{
	mmio_write_32(PLLDIG_PLLCR(pll_addr), PLLDIG_PLLCR_PLLPD);
}

static int program_pll(const struct s32cc_pll *pll, uintptr_t pll_addr,
		       const struct s32cc_clk_drv *drv, uint32_t sclk_id,
		       unsigned long sclk_freq)
{
	uint32_t rdiv = 1, mfi, mfn;
	int ret;

	ret = get_pll_mfi_mfn(pll->vco_freq, sclk_freq, &mfi, &mfn);
	if (ret != 0) {
		return -EINVAL;
	}

	/* Disable ODIVs*/
	disable_odivs(pll_addr, pll->ndividers);

	/* Disable PLL */
	disable_pll_hw(pll_addr);

	/* Program PLLCLKMUX */
	mmio_write_32(PLLDIG_PLLCLKMUX(pll_addr), sclk_id);

	/* Program VCO */
	mmio_clrsetbits_32(PLLDIG_PLLDV(pll_addr),
			   PLLDIG_PLLDV_RDIV_MASK | PLLDIG_PLLDV_MFI_MASK,
			   PLLDIG_PLLDV_RDIV_SET(rdiv) | PLLDIG_PLLDV_MFI(mfi));

	mmio_write_32(PLLDIG_PLLFD(pll_addr),
		      PLLDIG_PLLFD_MFN_SET(mfn) | PLLDIG_PLLFD_SMDEN);

	enable_pll_hw(pll_addr);

	return ret;
}

static int enable_pll(const struct s32cc_clk_obj *module,
		      const struct s32cc_clk_drv *drv,
		      unsigned int *depth)
{
	const struct s32cc_pll *pll = s32cc_obj2pll(module);
	const struct s32cc_clkmux *mux;
	uintptr_t pll_addr = UL(0x0);
	unsigned long sclk_freq;
	uint32_t sclk_id;
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	mux = get_pll_mux(pll);
	if (mux == NULL) {
		return -EINVAL;
	}

	if (pll->instance != mux->module) {
		ERROR("MUX type is not in sync with PLL ID\n");
		return -EINVAL;
	}

	ret = get_base_addr(pll->instance, drv, &pll_addr);
	if (ret != 0) {
		ERROR("Failed to detect PLL instance\n");
		return ret;
	}

	switch (mux->source_id) {
	case S32CC_CLK_FIRC:
		sclk_freq = 48U * MHZ;
		sclk_id = 0;
		break;
	case S32CC_CLK_FXOSC:
		sclk_freq = 40U * MHZ;
		sclk_id = 1;
		break;
	default:
		ERROR("Invalid source selection for PLL 0x%lx\n",
		      pll_addr);
		return -EINVAL;
	};

	return program_pll(pll, pll_addr, drv, sclk_id, sclk_freq);
}

static inline struct s32cc_pll *get_div_pll(const struct s32cc_pll_out_div *pdiv)
{
	const struct s32cc_clk_obj *parent;

	parent = pdiv->parent;
	if (parent == NULL) {
		ERROR("Failed to identify PLL divider's parent\n");
		return NULL;
	}

	if (parent->type != s32cc_pll_t) {
		ERROR("The parent of the divider is not a PLL instance\n");
		return NULL;
	}

	return s32cc_obj2pll(parent);
}

static void config_pll_out_div(uintptr_t pll_addr, uint32_t div_index, uint32_t dc)
{
	uint32_t pllodiv;
	uint32_t pdiv;

	pllodiv = mmio_read_32(PLLDIG_PLLODIV(pll_addr, div_index));
	pdiv = PLLDIG_PLLODIV_DIV(pllodiv);

	if (((pdiv + 1U) == dc) && ((pllodiv & PLLDIG_PLLODIV_DE) != 0U)) {
		return;
	}

	if ((pllodiv & PLLDIG_PLLODIV_DE) != 0U) {
		disable_odiv(pll_addr, div_index);
	}

	pllodiv = PLLDIG_PLLODIV_DIV_SET(dc - 1U);
	mmio_write_32(PLLDIG_PLLODIV(pll_addr, div_index), pllodiv);

	enable_odiv(pll_addr, div_index);
}

static int enable_pll_div(const struct s32cc_clk_obj *module,
			  const struct s32cc_clk_drv *drv,
			  unsigned int *depth)
{
	const struct s32cc_pll_out_div *pdiv = s32cc_obj2plldiv(module);
	uintptr_t pll_addr = 0x0ULL;
	const struct s32cc_pll *pll;
	uint32_t dc;
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	pll = get_div_pll(pdiv);
	if (pll == NULL) {
		ERROR("The parent of the PLL DIV is invalid\n");
		return 0;
	}

	ret = get_base_addr(pll->instance, drv, &pll_addr);
	if (ret != 0) {
		ERROR("Failed to detect PLL instance\n");
		return -EINVAL;
	}

	dc = (uint32_t)(pll->vco_freq / pdiv->freq);

	config_pll_out_div(pll_addr, pdiv->index, dc);

	return 0;
}

static int enable_module(const struct s32cc_clk_obj *module, unsigned int *depth)
{
	const struct s32cc_clk_drv *drv = get_drv();
	int ret = 0;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	if (drv == NULL) {
		return -EINVAL;
	}

	switch (module->type) {
	case s32cc_osc_t:
		ret = enable_osc(module, drv, depth);
		break;
	case s32cc_clk_t:
		ret = enable_clk_module(module, drv, depth);
		break;
	case s32cc_pll_t:
		ret = enable_pll(module, drv, depth);
		break;
	case s32cc_pll_out_div_t:
		ret = enable_pll_div(module, drv, depth);
		break;
	case s32cc_clkmux_t:
		ret = -ENOTSUP;
		break;
	case s32cc_shared_clkmux_t:
		ret = -ENOTSUP;
		break;
	case s32cc_fixed_div_t:
		ret = -ENOTSUP;
		break;
	default:
		ret = -EINVAL;
		break;
	}

	return ret;
}

static int s32cc_clk_enable(unsigned long id)
{
	unsigned int depth = MAX_STACK_DEPTH;
	const struct s32cc_clk *clk;

	clk = s32cc_get_arch_clk(id);
	if (clk == NULL) {
		return -EINVAL;
	}

	return enable_module(&clk->desc, &depth);
}

static void s32cc_clk_disable(unsigned long id)
{
}

static bool s32cc_clk_is_enabled(unsigned long id)
{
	return false;
}

static unsigned long s32cc_clk_get_rate(unsigned long id)
{
	return 0;
}

static int set_module_rate(const struct s32cc_clk_obj *module,
			   unsigned long rate, unsigned long *orate,
			   unsigned int *depth);

static int set_osc_freq(const struct s32cc_clk_obj *module, unsigned long rate,
			unsigned long *orate, unsigned int *depth)
{
	struct s32cc_osc *osc = s32cc_obj2osc(module);
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	if ((osc->freq != 0UL) && (rate != osc->freq)) {
		ERROR("Already initialized oscillator. freq = %lu\n",
		      osc->freq);
		return -EINVAL;
	}

	osc->freq = rate;
	*orate = osc->freq;

	return 0;
}

static int set_clk_freq(const struct s32cc_clk_obj *module, unsigned long rate,
			unsigned long *orate, unsigned int *depth)
{
	const struct s32cc_clk *clk = s32cc_obj2clk(module);
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	if ((clk->min_freq != 0UL) && (clk->max_freq != 0UL) &&
	    ((rate < clk->min_freq) || (rate > clk->max_freq))) {
		ERROR("%lu frequency is out of the allowed range: [%lu:%lu]\n",
		      rate, clk->min_freq, clk->max_freq);
		return -EINVAL;
	}

	if (clk->module != NULL) {
		return set_module_rate(clk->module, rate, orate, depth);
	}

	if (clk->pclock != NULL) {
		return set_clk_freq(&clk->pclock->desc, rate, orate, depth);
	}

	return -EINVAL;
}

static int set_pll_freq(const struct s32cc_clk_obj *module, unsigned long rate,
			unsigned long *orate, unsigned int *depth)
{
	struct s32cc_pll *pll = s32cc_obj2pll(module);
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	if ((pll->vco_freq != 0UL) && (pll->vco_freq != rate)) {
		ERROR("PLL frequency was already set\n");
		return -EINVAL;
	}

	pll->vco_freq = rate;
	*orate = pll->vco_freq;

	return 0;
}

static int set_pll_div_freq(const struct s32cc_clk_obj *module, unsigned long rate,
			    unsigned long *orate, unsigned int *depth)
{
	struct s32cc_pll_out_div *pdiv = s32cc_obj2plldiv(module);
	const struct s32cc_pll *pll;
	unsigned long prate, dc;
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	if (pdiv->parent == NULL) {
		ERROR("Failed to identify PLL divider's parent\n");
		return -EINVAL;
	}

	pll = s32cc_obj2pll(pdiv->parent);
	if (pll == NULL) {
		ERROR("The parent of the PLL DIV is invalid\n");
		return -EINVAL;
	}

	prate = pll->vco_freq;

	/**
	 * The PLL is not initialized yet, so let's take a risk
	 * and accept the proposed rate.
	 */
	if (prate == 0UL) {
		pdiv->freq = rate;
		*orate = rate;
		return 0;
	}

	/* Decline in case the rate cannot fit PLL's requirements. */
	dc = prate / rate;
	if ((prate / dc) != rate) {
		return -EINVAL;
	}

	pdiv->freq = rate;
	*orate = pdiv->freq;

	return 0;
}

static int set_fixed_div_freq(const struct s32cc_clk_obj *module, unsigned long rate,
			      unsigned long *orate, unsigned int *depth)
{
	const struct s32cc_fixed_div *fdiv = s32cc_obj2fixeddiv(module);
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	if (fdiv->parent == NULL) {
		ERROR("The divider doesn't have a valid parent\b");
		return -EINVAL;
	}

	ret = set_module_rate(fdiv->parent, rate * fdiv->rate_div, orate, depth);

	/* Update the output rate based on the parent's rate */
	*orate /= fdiv->rate_div;

	return ret;
}

static int set_mux_freq(const struct s32cc_clk_obj *module, unsigned long rate,
			unsigned long *orate, unsigned int *depth)
{
	const struct s32cc_clkmux *mux = s32cc_obj2clkmux(module);
	const struct s32cc_clk *clk = s32cc_get_arch_clk(mux->source_id);
	int ret;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	if (clk == NULL) {
		ERROR("Mux (id:%" PRIu8 ") without a valid source (%lu)\n",
		      mux->index, mux->source_id);
		return -EINVAL;
	}

	return set_module_rate(&clk->desc, rate, orate, depth);
}

static int set_module_rate(const struct s32cc_clk_obj *module,
			   unsigned long rate, unsigned long *orate,
			   unsigned int *depth)
{
	int ret = 0;

	ret = update_stack_depth(depth);
	if (ret != 0) {
		return ret;
	}

	switch (module->type) {
	case s32cc_clk_t:
		ret = set_clk_freq(module, rate, orate, depth);
		break;
	case s32cc_osc_t:
		ret = set_osc_freq(module, rate, orate, depth);
		break;
	case s32cc_pll_t:
		ret = set_pll_freq(module, rate, orate, depth);
		break;
	case s32cc_pll_out_div_t:
		ret = set_pll_div_freq(module, rate, orate, depth);
		break;
	case s32cc_fixed_div_t:
		ret = set_fixed_div_freq(module, rate, orate, depth);
		break;
	case s32cc_clkmux_t:
		ret = set_mux_freq(module, rate, orate, depth);
		break;
	case s32cc_shared_clkmux_t:
		ret = set_mux_freq(module, rate, orate, depth);
		break;
	default:
		ret = -EINVAL;
		break;
	}

	return ret;
}

static int s32cc_clk_set_rate(unsigned long id, unsigned long rate,
			      unsigned long *orate)
{
	unsigned int depth = MAX_STACK_DEPTH;
	const struct s32cc_clk *clk;
	int ret;

	clk = s32cc_get_arch_clk(id);
	if (clk == NULL) {
		return -EINVAL;
	}

	ret = set_module_rate(&clk->desc, rate, orate, &depth);
	if (ret != 0) {
		ERROR("Failed to set frequency (%lu MHz) for clock %lu\n",
		      rate, id);
	}

	return ret;
}

static int s32cc_clk_get_parent(unsigned long id)
{
	return -ENOTSUP;
}

static int s32cc_clk_set_parent(unsigned long id, unsigned long parent_id)
{
	const struct s32cc_clk *parent;
	const struct s32cc_clk *clk;
	bool valid_source = false;
	struct s32cc_clkmux *mux;
	uint8_t i;

	clk = s32cc_get_arch_clk(id);
	if (clk == NULL) {
		return -EINVAL;
	}

	parent = s32cc_get_arch_clk(parent_id);
	if (parent == NULL) {
		return -EINVAL;
	}

	if (!is_s32cc_clk_mux(clk)) {
		ERROR("Clock %lu is not a mux\n", id);
		return -EINVAL;
	}

	mux = s32cc_clk2mux(clk);
	if (mux == NULL) {
		ERROR("Failed to cast clock %lu to clock mux\n", id);
		return -EINVAL;
	}

	for (i = 0; i < mux->nclks; i++) {
		if (mux->clkids[i] == parent_id) {
			valid_source = true;
			break;
		}
	}

	if (!valid_source) {
		ERROR("Clock %lu is not a valid clock for mux %lu\n",
		      parent_id, id);
		return -EINVAL;
	}

	mux->source_id = parent_id;

	return 0;
}

void s32cc_clk_register_drv(void)
{
	static const struct clk_ops s32cc_clk_ops = {
		.enable		= s32cc_clk_enable,
		.disable	= s32cc_clk_disable,
		.is_enabled	= s32cc_clk_is_enabled,
		.get_rate	= s32cc_clk_get_rate,
		.set_rate	= s32cc_clk_set_rate,
		.get_parent	= s32cc_clk_get_parent,
		.set_parent	= s32cc_clk_set_parent,
	};

	clk_register(&s32cc_clk_ops);
}