// SPDX-License-Identifier: (GPL-2.0 OR BSD-3-Clause)
/*
 * Copyright (c) 2020 MediaTek Inc.
 */

#include <linux/of_device.h>
#include <linux/io.h>
#include <linux/spinlock.h>
#include <linux/mutex.h>
#include <linux/interrupt.h>
#include <linux/completion.h>
#include <linux/slab.h>
#include <linux/mm.h>
#include <linux/debugfs.h>
#include <linux/version.h>
#include <linux/dma-direct.h>
#include <linux/sched.h>
#include <linux/sched/clock.h>
#include <asm/atomic.h>

#include <common/mdla_device.h>
#include <common/mdla_cmd_proc.h>
#include <common/mdla_power_ctrl.h>
#include <common/mdla_ioctl.h>
#include <common/mdla_scheduler.h>
#include <utilities/mdla_debug.h>
#include <utilities/mdla_profile.h>
#include <utilities/mdla_util.h>
#include <utilities/mdla_trace.h>
#include <interface/mdla_cmd_data_v2_0.h>
#include <platform/mdla_plat_api.h>
#include <platform/v2_0/mdla_plat_teec.h>
#include <platform/v2_0/mdla_plat_model.h>

static DEFINE_MUTEX(model_lock);
static atomic_t model_load_en;
static atomic_t model_param_en;
static u32 current_ctx;

static LIST_HEAD(model_param_list);
static LIST_HEAD(model_load_list);
static LIST_HEAD(model_unload_list);

static inline bool model_context_vaild(struct model_info *model)
{
	if (atomic_read(&model_load_en) == 0)
		return true;

	if (model->model_ctx >= current_ctx)
		return true;

	mdla_err("%s(): model(%x) model_ctx(%x) < current_ctx(%x)\n",
		__func__, model->model_idx, model->model_ctx, current_ctx);

	return false;
}

bool model_check_runable(void *data)
{
	struct model_info *model = data;

	if (!model)
		return false;

	if (atomic_read(&model_param_en) == 0)
		return model_context_vaild(model);

	if (model->param)
		return true;

	mdla_err("%s(): model:%x doesn't have update info\n",
		__func__, model->model_idx);

	return false;
}

int mdla_plat_model_reset(void)
{
	current_ctx++;

	return 0;
}

static struct model_info *model_create(u32 model_idx, u32 owner)
{
	struct model_info *p;

	p = vzalloc(sizeof(struct model_info));
	if (!p)
		return NULL;

	init_completion(&p->cmplt);
	kref_init(&p->refcount);

	p->model_idx = model_idx;
	p->model_ctx = current_ctx;
	p->owner = owner;

	INIT_LIST_HEAD(&p->mdata);
	INIT_LIST_HEAD(&p->mdata_del);
	INIT_LIST_HEAD(&p->pmi_list);

	list_add(&p->mt_list, &model_load_list);

	return p;
}

static void model_destroy(struct kref *kref)
{
	struct model_info *m =
			container_of(kref, struct model_info, refcount);
	struct list_head *tmp, *next;
	struct model_data *d = NULL;

	list_del(&m->mt_list);

	list_for_each_safe(tmp, next, &m->mdata) {
		d = list_entry(tmp, struct model_data, list);
		list_del(&d->list);
		vfree(d);
	}

	list_for_each_safe(tmp, next, &m->mdata_del) {
		d = list_entry(tmp, struct model_data, list);
		list_del(&d->list);
		vfree(d);
	}

	vfree(m);
}

static void model_delete(struct model_info *m)
{
	list_move(&m->mt_list, &model_unload_list);
	kref_put(&m->refcount, model_destroy);
}

static struct model_info *model_find(u32 model_idx, u32 create)
{
	struct list_head *tmp, *next;
	struct model_info *p = NULL;

	list_for_each_safe(tmp, next, &model_load_list) {
		p = list_entry(tmp, struct model_info, mt_list);

		if (p->model_idx == model_idx)
			break;

		p = NULL;
	}

	if (create) {
		/* reclaim model if it's outdue and alloc internal */
		if (p && p->model_ctx < current_ctx) {
			if (p->owner == 0) {
				model_delete(p);
				p = NULL;
			}
		}

		if (p == NULL && atomic_read(&model_load_en) == 0)
			p = model_create(model_idx, 0);
	}

	return p;
}

#define seq_redirect(s, ...)			\
do {						\
	if (s)					\
		seq_printf(s, __VA_ARGS__);	\
	else					\
		redirect_output(__VA_ARGS__);	\
} while (0)

static void dump_model_data(char *prefix, struct seq_file *s, struct model_info *p)
{
	struct list_head *tmp, *next;
	struct model_data *d = NULL;
	u32 count = 0;

	list_for_each_safe(tmp, next, &p->mdata) {
		d = list_entry(tmp, struct model_data, list);
		seq_redirect(s, "%s data(%x) size(%x), kref(%x)\n",
			prefix, count, d->size, kref_read(&d->refcount));
		count++;
	}

	count = 0;
	list_for_each_safe(tmp, next, &p->mdata_del) {
		d = list_entry(tmp, struct model_data, list);
		seq_redirect(s, "%s del_data(%x) size(%x), kref(%x)\n",
			prefix, count, d->size, kref_read(&d->refcount));
		count++;
	}

}

void mdla_plat_model_dump_info(struct seq_file *s)
{
	struct list_head *tmp, *next;
	struct list_head *m_tmp, *m_next;
	struct model_param *mp = NULL;
	struct model_info *p = NULL;

	mutex_lock(&model_lock);

	seq_redirect(s, "---- dump MDLA model info ----\n");
	seq_redirect(s, "model_load_en: %x\n", atomic_read(&model_load_en));
	seq_redirect(s, "model_param_en: %x\n", atomic_read(&model_param_en));
	seq_redirect(s, "current_ctx: %x\n", current_ctx);
	seq_redirect(s, "[model param]\n");
	list_for_each_safe(tmp, next, &model_param_list) {
		mp = list_entry(tmp, struct model_param, mp_list);

		seq_redirect(s, "|- pipe_id:%x, num_model:%x,  num_buf:%x\n",
			mp->pipe_id, mp->num_model, mp->num_buf);

		list_for_each_safe(m_tmp, m_next, &mp->mi_list) {
			p = list_entry(m_tmp, struct model_info, pmi_list);

			seq_redirect(s, "\t|- model:%x ctx:%x, ref:%x\n",
				p->model_idx, p->model_ctx,
				kref_read(&p->refcount));
		}
	}

	seq_redirect(s, "[model open]\n");
	list_for_each_safe(tmp, next, &model_load_list) {
		p = list_entry(tmp, struct model_info, mt_list);

		seq_redirect(s, "|- model:%x, ctx:%x, owner:%x, ref:%x\n",
			p->model_idx, p->model_ctx,
			p->owner, kref_read(&p->refcount));
		dump_model_data("\t|- ", s, p);
	}

	seq_redirect(s, "[model delete]\n");
	list_for_each_safe(tmp, next, &model_unload_list) {
		p = list_entry(tmp, struct model_info, mt_list);

		seq_redirect(s, "|- model:%x, ctx:%x, owner:%x, ref:%x\n",
			p->model_idx, p->model_ctx,
			p->owner, kref_read(&p->refcount));
		dump_model_data("\t |- ", s, p);
	}

	mutex_unlock(&model_lock);
}

static void model_data_delete(struct kref *kref)
{
	struct model_data *d =
			container_of(kref, struct model_data, refcount);

	list_del(&d->list);
	vfree(d);
}

// NOTE: decrease refcount of current data and insert new data
int mdla_plat_model_set_data(u32 model_idx, u32 size, void *data)
{
	struct model_info *p;
	struct model_data *md = NULL;
	struct model_data *tmp;
	u32 md_size;

	mutex_lock(&model_lock);

	p = model_find(model_idx, 1);
	if (!p || !model_context_vaild(p)) {
		mutex_unlock(&model_lock);
		return -EFAULT;
	}

	if (size) {
		md_size = sizeof(struct model_data) + size;
		md = vzalloc(md_size);
		if (md == NULL) {
			mutex_unlock(&model_lock);
			return -ENOMEM;
		}

		kref_init(&md->refcount);
		INIT_LIST_HEAD(&md->list);
		md->size = size;
		md->model = p;
		memcpy(md->data, data, size);
	}

	if (!list_empty(&p->mdata)) {
		tmp = list_first_entry(&p->mdata,
			struct model_data, list);
		list_move(&tmp->list, &p->mdata_del);
		kref_put(&tmp->refcount, model_data_delete);
	}

	if (md)
		list_add(&md->list, &p->mdata);

	mutex_unlock(&model_lock);

	return 0;
}

// NOTE: increase refcount of model and model_data
int mdla_plat_model_get_data(u32 model_idx, void **data)
{
	struct model_info *p;
	struct model_data *md = NULL;

	mutex_lock(&model_lock);

	p = model_find(model_idx, 1);
	if (!p || !model_context_vaild(p)) {
		mutex_unlock(&model_lock);
		return -EFAULT;
	}

	if (!list_empty(&p->mdata)) {
		md = list_first_entry(&p->mdata,
			struct model_data, list);
		kref_get(&p->refcount);
		kref_get(&md->refcount);
	}

	*data = (md) ? md->data : NULL;

	mutex_unlock(&model_lock);

	return 0;
}

// NOTE: decrease refcount of model and model_data
int mdla_plat_model_put_data(u32 model_idx, void *data)
{
	struct model_info *model;
	struct model_data *mdata;

	if (data == NULL)
		return 0;

	mdata = container_of(data, struct model_data, data);
	model = mdata->model;

	mutex_lock(&model_lock);

	kref_put(&mdata->refcount, model_data_delete);
	kref_put(&model->refcount, model_destroy);

	mutex_unlock(&model_lock);

	return 0;
}

int mdla_plat_model_get(u32 model_idx, void **model)
{
	struct model_info *tmp;

	mutex_lock(&model_lock);

	tmp = model_find(model_idx, 1);
	if (!tmp || !model_context_vaild(tmp)) {
		mutex_unlock(&model_lock);
		return -EFAULT;
	}

	kref_get(&tmp->refcount);
	*model = tmp;

	mutex_unlock(&model_lock);

	return 0;
}

int mdla_plat_model_put(void *model)
{
	struct model_info *tmp = (struct model_info *)model;

	mutex_lock(&model_lock);

	kref_put(&tmp->refcount, model_destroy);

	mutex_unlock(&model_lock);

	return 0;
}

int mdla_plat_model_load(void *data)
{
	struct ioctl_load_model_svp *model =
		(struct ioctl_load_model_svp *)data;
	struct model_info *p;

	mutex_lock(&model_lock);

	atomic_set(&model_load_en, 1);

	p = model_find(model->model, 0);
	if (p) {
		mdla_err("%s(): model(%x) not unload, model_ctx:%x\n",
				__func__, model->model, p->model_ctx);
		model_delete(p);
	}

	p = model_create(model->model, 1);
	if (p) {
		p->cmdbuf_addr = model->buf_addr;
		p->cmdbuf_size = model->buf_size;
		p->model_offset = 0xFFFFFFFF;
	}
	mutex_unlock(&model_lock);

	return (p) ? 0 : -ENOMEM;
}

int mdla_plat_model_unload(void *data)
{
	struct ioctl_load_model_svp *model =
		(struct ioctl_load_model_svp *)data;
	struct model_info *p;

	mutex_lock(&model_lock);

	p = model_find(model->model, 0);
	if (!p) {
		mdla_err("%s(): model(%x) was not loaded\n",
				__func__, model->model);
		mutex_unlock(&model_lock);
		return -EFAULT;
	}

	if (kref_read(&p->refcount) > 1)
		mdla_err("%s(): model(%x) unload, kref(%x) not 0\n",
			__func__, model->model, kref_read(&p->refcount));

	model_delete(p);

	mutex_unlock(&model_lock);

	return 0;
}

static bool model_param_data_match(struct model_param *p, struct ioctl_update_model_svp *data)
{
	int i, j, cnt;

	if (p->num_buf != data->num_buf)
		return false;

	cnt = 0;

	for (i = 0; i < p->num_buf; i++) {
		for (j = 0; j < data->num_buf; j++) {
			if (memcmp(&p->buf_info[i], &data->buf_list[j],
					sizeof(struct mdla_svp_buf)) == 0) {
				cnt++;
				break;
			}
		}
	}

	return cnt == p->num_buf;
}

static void model_param_destroy(struct kref *kref)
{
	struct model_param *p =
			container_of(kref, struct model_param, refcount);

	list_del(&p->mp_list);
	vfree(p);
}

static void model_param_detach_model(struct model_info *m)
{
	struct model_param *p = m->param;

	if (!p)
		return;

	m->param = NULL;
	list_del(&m->pmi_list);
	p->num_model--;

	kref_put(&m->refcount, model_destroy);
	kref_put(&p->refcount, model_param_destroy);
}

static void model_param_attach_model(struct model_param *p, struct model_info *m)
{
	kref_get(&p->refcount);
	kref_get(&m->refcount);

	if (m->param) {
		mdla_err("%s(): pipe:%x->%x model:%x update but not reset\n",
			__func__, m->param->pipe_id, p->pipe_id, m->model_idx);

		model_param_detach_model(m);
	}

	m->param = p;
	list_add_tail(&m->pmi_list, &p->mi_list);
}

static void model_param_delete(struct model_param *p)
{
	struct list_head *tmp, *next;
	struct model_info *m = NULL;

	list_for_each_safe(tmp, next, &p->mi_list) {
		m = list_entry(tmp, struct model_info, pmi_list);
		model_param_detach_model(m);
	}

	kref_put(&p->refcount, model_param_destroy);
}

static struct model_param *model_param_create(void *user, struct ioctl_update_model_svp *data)
{
	struct model_param *p;

	p = vzalloc(sizeof(struct model_param));
	if (!p)
		return NULL;

	p->user = user;
	p->pipe_id = data->pipe_id;
	p->num_model = data->num_model;

	kref_init(&p->refcount);
	INIT_LIST_HEAD(&p->mp_list);
	INIT_LIST_HEAD(&p->mi_list);

	p->num_buf = data->num_buf;
	memcpy(p->buf_info, data->buf_list, p->num_buf * sizeof(struct mdla_svp_buf));

	list_add(&p->mp_list, &model_param_list);

	return p;
}

int mdla_plat_model_param_update(void *user, void *data)
{
	struct ioctl_update_model_svp *param = data;
	struct model_param *p = NULL;
	struct model_info *m;
	struct list_head *tmp, *next;
	int i;

	atomic_set(&model_param_en, 1);

	mutex_lock(&model_lock);

	/* find param with same data then merge it */
	list_for_each_safe(tmp, next, &model_param_list) {
		p = list_entry(tmp, struct model_param, mp_list);

		if (p->user == user && p->pipe_id == param->pipe_id) {
			if (model_param_data_match(p, param) == true) {
				mdla_err("%s(): pipe:%x, update with same param\n",
					__func__, p->pipe_id);
				break;
			}
		}

		p = NULL;
	}

	if (!p) {
		p = model_param_create(user, param);
		if (!p) {
			mdla_err("%s(): create param pipe:%x fail\n",
					__func__, param->pipe_id);
			mutex_unlock(&model_lock);
			return -EFAULT;
		}
			mdla_err("%s(): new param pipe:%x\n",
					__func__, param->pipe_id);
	}

	for (i = 0; i < param->num_model; i++) {
		m = model_find(param->model_list[i], 0);

		/* keep insert param event model is invalid for reset on release */
		if (!m || !model_context_vaild(m) || m->owner == 0)
			mdla_err("%s(): pipe:%x, model:%x not valid\n",
					__func__, param->pipe_id, param->model_list[i]);

		if (m)
			model_param_attach_model(p, m);
	}

	mutex_unlock(&model_lock);

	return 0;
}

int mdla_plat_model_param_reset(void *user, void *data)
{
	struct ioctl_update_model_svp *param = data;
	struct model_param *p;
	struct model_info *m;
	struct list_head *tmp, *next;
	u32 count = 0;
	int i;

	atomic_set(&model_param_en, 1);

	mutex_lock(&model_lock);

	for (i = 0; i < param->num_model; i++) {
		m = model_find(param->model_list[i], 0);
		if (!m || !m->param) {
			mdla_err("%s(): pipe:%x, model:%x didn't load/update\n",
					__func__, param->pipe_id, param->model_list[i]);
			continue;
		}

		p = m->param;

		if (p->user != user || p->pipe_id != param->pipe_id) {
			mdla_err("%s(): pipe:%x, model:%x user/pipe mismatch\n",
					__func__, param->pipe_id, param->model_list[i]);
			continue;
		}

		if (model_param_data_match(p, param) == false)
			mdla_err("%s(): pipe:%x, model:%x svp_buf mismatch\n",
				__func__, param->pipe_id, param->model_list[i]);

		model_param_detach_model(m);
		count++;
	}

	list_for_each_safe(tmp, next, &model_param_list) {
		p = list_entry(tmp, struct model_param, mp_list);

		if (p->user == user && p->pipe_id == param->pipe_id) {
			if (model_param_data_match(p, param) == true) {

				kref_put(&p->refcount, model_param_destroy);

				mdla_err("%s(): pipe:%x detach (%x/%x) model\n",
					__func__, param->pipe_id, count, param->num_model);
				break;
			}
		}

	}

	mutex_unlock(&model_lock);

	return 0;
}

static int model_teec_param_reset(struct model_param *p)
{
	struct list_head *tmp, *next;
	struct model_info *m;
	u32 count;
	u32 model_list[MAX_AI_MODEL];
	struct mdla_teec_param_update data;

	count = 0;

	list_for_each_safe(tmp, next, &p->mi_list) {
		m = list_entry(tmp, struct model_info, pmi_list);
		model_list[count++] = m->model_idx;
	}

	data.pipe_id = p->pipe_id;
	data.num_model = count;
	data.model_list = model_list;
	data.num_buf = p->num_buf;
	data.buf_list = p->buf_info;

	return mdla_plat_teec_param_reset(&data);
}

static int model_teec_model_unload(struct model_info *m)
{
	struct mdla_teec_model_load data;


	data.model_idx = m->model_idx;
	data.buf_addr = m->cmdbuf_addr;
	data.buf_size = m->cmdbuf_size;
	data.buf_offset = m->model_offset;

	return mdla_plat_teec_unload_model(&data);
}

int mdla_plat_model_param_release(void *user)
{
	struct model_param *p;
	struct model_info *m;
	struct list_head *tmp, *next;
	struct list_head *m_tmp, *m_next;

	mutex_lock(&model_lock);
	list_for_each_safe(tmp, next, &model_param_list) {
		p = list_entry(tmp, struct model_param, mp_list);
		if (p->user == user) {
			mdla_err("%s(): pipe:%x leak, reset...\n",
				__func__, p->pipe_id);

			model_teec_param_reset(p);

			list_for_each_safe(m_tmp, m_next, &p->mi_list) {
				m = list_entry(m_tmp, struct model_info, pmi_list);

				mdla_err("%s(): unload model:%x buf:%llx\n",
					__func__, m->model_idx, m->cmdbuf_addr);

				model_teec_model_unload(m);
			}

			model_param_delete(p);
		}

	}
	mutex_unlock(&model_lock);

	return 0;
}
