// SPDX-License-Identifier: GPL-2.0
//
// Copyright (c) 2021 MediaTek Inc.

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/dma-buf.h>
#include <linux/types.h>
#include <linux/fs.h>
#include <linux/string.h>
#include <linux/proc_fs.h>
#include <linux/seq_file.h>
#include <linux/sizes.h>
#include <linux/slab.h>
#include <linux/iommu.h>
#include <linux/platform_device.h>
#include <linux/version.h>

#include <linux/kthread.h>
#include <linux/delay.h>
#if LINUX_VERSION_CODE <= KERNEL_VERSION(5, 10, 0)
#include <linux/ion.h>
#endif
#include <linux/freezer.h>
#include "mtk_iommu_dbg.h"

#include <linux/dma-heap.h>
#include <uapi/linux/dma-heap.h>

#undef pr_fmt
#define pr_fmt(fmt) "IOMMUTORTURE: " fmt

#define ALLOC_SIZE 0x8000
#define WRITE_VALUE 0x55

int alloc_free_times;

static atomic_t lock_stop = ATOMIC_INIT(0);
static struct device* dev = NULL;
static const struct of_device_id mtk_test_of_ids[] = {
	{ .compatible = "mediatek,dtv-iommu-test",},
	{},
};

static int mtk_iommu_test_probe(struct platform_device* pdev)
{
	dev = &pdev->dev;

	return 0;
}

static int mtk_iommu_test_remove(struct platform_device *pdev)
{
	iommu_fwspec_free(&pdev->dev);
	return 0;
}

static struct platform_driver mtk_iommu_test_driver = {
	.probe  = mtk_iommu_test_probe,
	.remove = mtk_iommu_test_remove,
	.driver = {
        	.name = "mtk-iommu-test",
		.of_match_table = mtk_test_of_ids,
	}
};

static int mtk_iommu_test_init(void)
{
	return platform_driver_register(&mtk_iommu_test_driver);
}

static void mtk_iommu_test_exit(void)
{
	platform_driver_unregister(&mtk_iommu_test_driver);
}

static int alloc_free_thread(void *data)
{
	struct dma_heap *dmaheap;
	struct dma_buf *dmabuf;
	struct dma_buf_attachment *attach;
#if LINUX_VERSION_CODE >= KERNEL_VERSION(5,15,0)
	struct dma_buf_map map = { };
#endif
	char *heap_name = atomic_read(&lock_stop) & 0xF ? "mtk_vdec_fb_uncached" : "system";
	struct sg_table *sgt;
	void *va;
	size_t size = ALLOC_SIZE;
	int ret = -EINVAL;

	dmaheap = dma_heap_find(heap_name);
	if (!dmaheap) {
		pr_emerg("find [%s] fail\n", heap_name);
		return -EINVAL;
	}

	dmabuf = dma_heap_buffer_alloc(dmaheap, size,
			DMA_HEAP_VALID_FD_FLAGS,
			DMA_HEAP_VALID_HEAP_FLAGS);
	if (IS_ERR_OR_NULL(dmabuf)) {
		pr_emerg("%s: [%s] alloc dma_buf failed ret: %ld\n", __func__,
				heap_name, PTR_ERR(dmabuf));
		goto out;
	}

	attach = dma_buf_attach(dmabuf, dev);
	if (IS_ERR(attach)) {
		goto put;
	}

	sgt = dma_buf_map_attachment(attach, DMA_BIDIRECTIONAL);
	if (IS_ERR(sgt)) {
		goto free_attach;
	}

	// begin cpu access
	ret = dma_buf_begin_cpu_access(dmabuf, DMA_BIDIRECTIONAL);
	if (ret < 0) {
		pr_emerg("%s: dma_buf_begin_cpu_access fail\n", __func__);
		goto free_map;
	}

	// vmap to kernel va
#if LINUX_VERSION_CODE >= KERNEL_VERSION(5,15,0)
	ret = dma_buf_vmap(dmabuf, &map);
	if (ret) {
#else
	va = dma_buf_vmap(dmabuf);
	if (IS_ERR_OR_NULL(va)) {
#endif
		dma_buf_end_cpu_access(dmabuf, DMA_BIDIRECTIONAL);
		pr_emerg("%s: dma_buf_vmap fail\n", __func__);
		goto free_map;
	}

#if LINUX_VERSION_CODE >= KERNEL_VERSION(5,15,0)
	va = map.vaddr;
#endif
	memset(va, WRITE_VALUE, size);

	// vunmap kernel va
#if LINUX_VERSION_CODE >= KERNEL_VERSION(5,15,0)
	dma_buf_vunmap(dmabuf, &map);
#else
	dma_buf_vunmap(dmabuf, va);
#endif

	// end cpu access
	dma_buf_end_cpu_access(dmabuf, DMA_BIDIRECTIONAL);
	if (ret < 0) {
		pr_emerg("%s: dma_buf_end_cpu_access fail\n", __func__);
		goto free_map;
	}

free_map:
	dma_buf_unmap_attachment(attach, sgt, DMA_BIDIRECTIONAL);
free_attach:
	dma_buf_detach(dmabuf, attach);
put:
	dma_buf_put(dmabuf);
out:
	dma_heap_put(dmaheap);

	atomic_inc(&lock_stop);
	return ret;
}

static struct task_struct *af_tsk;

int alloc_free_start(unsigned long times)
{
	int i;
	alloc_free_times = 0;
	atomic_set(&lock_stop, 0);

	mtk_iommu_test_init();
	pr_crit("torture start...\n");
	for (i = 0; i < times; i++) {
		alloc_free_times++;
		kthread_run(alloc_free_thread, NULL, "iommu_stress_%d", i);
	}
	while (atomic_read(&lock_stop) < times);
	pr_crit("torture finish...\n");
	mtk_iommu_test_exit();
	return 0;
}

int alloc_free_stop(void)
{
	if (!af_tsk) {
		pr_err("the stress thread is not start yet");
		return -EINVAL;
	}
	kthread_stop(af_tsk);
	af_tsk = NULL;

	return 0;
}

long get_alloc_free_cnt(void)
{
	return alloc_free_times;
}

static struct task_struct *lk_tsk;
static size_t leaker_buffer_size = 0x810000;
struct dma_buf *dmabuf;

static int leaker_thread(void *data)
{
	int i = 1, fd;
	struct dma_buf *db = (struct dma_buf *)data;

	pr_crit("%s start with dma 0x%lx\n", __func__, (unsigned long)data);

	fd = dma_buf_fd(db, O_CLOEXEC);
	if (fd < 0) {
		pr_crit("%s: get fd failed with %d\n",
			__func__, fd);
		return fd;
	}

	do {
		try_to_freeze();

		db = dma_buf_get(fd);

		if (i % 5 == 0)
			dma_buf_put(dmabuf);

		ssleep(3);
		i++;
		pr_info("%s: dma 0x%lx, file count %d\n",
			__func__, (unsigned long)dmabuf,
			(int)file_count(dmabuf->file));

	} while (!kthread_should_stop());

	pr_crit("%s exit\n", __func__);

	return 0;
}

int leaker_start(void)
{
	int rc;
	struct dma_buf *db;
	struct dma_heap *dmaheap;
	char *heap_name = "mtk_vdec_es";

	if (lk_tsk) {
		pr_err("the leaker thread is start already\n");
		return -EINVAL;
	}

	dmaheap = dma_heap_find(heap_name);
	if (!dmaheap) {
		pr_err("find [%s] fail\n", heap_name);
		return -EINVAL;
	}

	db = dma_heap_buffer_alloc(dmaheap, ALLOC_SIZE,
			DMA_HEAP_VALID_FD_FLAGS,
			DMA_HEAP_VALID_HEAP_FLAGS);
	if (IS_ERR_OR_NULL(db)) {
		pr_err("%s: [%s] alloc dma_buf failed ret: %ld\n", __func__,
				heap_name, PTR_ERR(db));
		dma_heap_put(dmaheap);
		return PTR_ERR(db);
	}

	// free
	dma_heap_put(dmaheap);
	dmabuf = db;

	pr_info("alloc iommu buffer with %zx bytes at 0x%lx, file count = %d\n",
		leaker_buffer_size, (unsigned long)db,
		(int)file_count(db->file));

	lk_tsk = kthread_run(leaker_thread, db, "iommu_leaker");
	if (IS_ERR(lk_tsk)) {
		rc = PTR_ERR(lk_tsk);
		pr_err("kthread create failed with %d\n", rc);
		lk_tsk = NULL;
		return rc;
	}

	return 0;
}

int leaker_stop(void)
{
	int i, fcnt;

	if (!lk_tsk) {
		pr_err("the leaker thread is not start yet");
		return -EINVAL;
	}

	if (!dmabuf) {
		pr_err("the dma buf is NULL\n");
		return -ENOMEM;
	}

	kthread_stop(lk_tsk);
	lk_tsk = NULL;
	fcnt = (int)file_count(dmabuf->file);

	pr_crit("total %d file count, put it!\n", fcnt);

	for (i = 0; i < fcnt; ++i)
		dma_buf_put(dmabuf);

	dmabuf = NULL;
	return 0;
}

struct dma_buf *get_leak_dmabuf(void)
{
	return dmabuf;
}
