#include <linux/module.h>
#include <linux/arm-smccc.h>
#include <linux/errno.h>
#include <linux/string.h>
#include <linux/of.h>
#include <linux/fs.h>
#include <linux/io.h>
#include <linux/proc_fs.h>
#include <linux/delay.h>
#include <linux/ctype.h>
#include <linux/slab.h>
#include <linux/compat.h>
#include <linux/kprobes.h>
#include <../drivers/tee/optee/optee_private.h>
#include <../drivers/tee/optee/optee_smc.h>
#include <linux/tee_drv.h>
#include <linux/tee.h>
#include <linux/sizes.h>
#include <linux/kprobes.h>
#include <asm/ptrace.h>
#include <linux/timekeeping.h>
#include <linux/seq_file.h>

static LIST_HEAD(invoke_record_head);
static DEFINE_SPINLOCK(invoke_list_lock);
ktime_t start_rec_time;
//static LIST_HEAD(open_sess_record_head);
//static DEFINE_SPINLOCK(open_sess_list_lock);

struct smc_record {
	pid_t pid;
	pid_t tgid;
	char comm[TASK_COMM_LEN];
	unsigned long cnt;
	ktime_t last_time;
	ktime_t last_cost_time;
	ktime_t sum_cost_time;
	uint32_t func;
	struct list_head node;
};

static struct smc_record *alloc_smc_node(pid_t pid, uint32_t func)
{
	struct smc_record *new;
	int ret;

	new = kzalloc(sizeof(struct smc_record), GFP_KERNEL);
	if (!new) {
		pr_crit("%s kzalloc fail\n", __func__);
		return NULL;
	}

	new->pid = pid;
	new->tgid = current->tgid;
	new->func = func;
	new->last_cost_time = 0;
	new->sum_cost_time = 0;
	ret = snprintf(new->comm, sizeof(new->comm), "%s", current->comm);
	if (ret < 0)
		pr_crit("%s copy comm happen error: %x\n", __func__, ret);
	new->cnt = 1;
	new->last_time = ktime_get_coarse();

	return new;
}

static void add_to_invoke_list(pid_t pid, uint32_t func)
{
	struct smc_record *tmp, *traversal, *new;
	bool found = false;
	unsigned long flags;

	spin_lock_irqsave(&invoke_list_lock, flags);
	list_for_each_entry_safe(traversal, tmp, &invoke_record_head, node) {
		// use pid as index
		if (traversal->pid == pid) {
			found = true;
			break;
		}
	}

	if (found) {
		traversal->cnt++;
		traversal->last_cost_time = 0;
		traversal->last_time = ktime_get_coarse();
		traversal->func = func;
		goto out;
	}

	new = alloc_smc_node(pid, func);
	if (!new) {
		pr_crit("%s %d\n", __func__, __LINE__);
		goto out;
	}

	INIT_LIST_HEAD(&new->node);
	list_add(&new->node, &invoke_record_head);
out:
	spin_unlock_irqrestore(&invoke_list_lock, flags);
	return;
}

static void modify_finish_time(pid_t pid)
{
	struct smc_record *tmp, *traversal;
	bool found = false;
	unsigned long flags;

	spin_lock_irqsave(&invoke_list_lock, flags);
	list_for_each_entry_safe(traversal, tmp, &invoke_record_head, node) {
		// use pid as index
		if (traversal->pid == pid) {
			found = true;
			break;
		}
	}

	if (found) {
		traversal->last_cost_time = ktime_get_coarse() - traversal->last_time;
		traversal->sum_cost_time += traversal->last_cost_time;
		spin_unlock_irqrestore(&invoke_list_lock, flags);
		return;
	}

	spin_unlock_irqrestore(&invoke_list_lock, flags);
	pr_debug("%s Something wrong...\n",__func__);
	return;
}

static int smc_hook(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct tee_ioctl_invoke_arg *arg;

	if (regs) {
		arg = (void*)regs_get_register(regs, 1);
	} else {
		pr_crit("%s %d ???\n", __func__, __LINE__);
		return 0;
	}

	add_to_invoke_list(current->pid, arg->func);
	return 0;
}

static int smc_ret_hook(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	modify_finish_time(current->pid);
	return 0;
}

static struct kretprobe probe_invoke = {
	.entry_handler  = smc_hook,
	.handler	= smc_ret_hook,
};

int smc_info_dump(struct seq_file *m, void *v)
{
	struct smc_record *tmp, *traversal;
	unsigned long flags;

	seq_printf(m, "Dump SMC History: Start to record at kernel time:%lld ms\n"
			,ktime_to_ms(start_rec_time));
	spin_lock_irqsave(&invoke_list_lock, flags);
	list_for_each_entry_safe(traversal, tmp, &invoke_record_head, node) {
		seq_printf(m, "%10s pid:%5d tgid:%5d, cnt:%5ld, func:0x%x last(%lld ms, cost:%lld ms) sum cost: %lld ms\n",
			traversal->comm, traversal->pid, traversal->tgid,
			traversal->cnt, traversal->func,
			ktime_to_ms(traversal->last_time),
			ktime_to_ms(traversal->last_cost_time),
			ktime_to_ms(traversal->sum_cost_time));
	}
	spin_unlock_irqrestore(&invoke_list_lock, flags);
	return 0;
}

int register_symbol(unsigned long addr)
{
	int ret;

	pr_err("[In] register_symbol addr: 0x%lx\n", addr);
	probe_invoke.kp.addr = (kprobe_opcode_t *)(void*)addr;

	ret = register_kretprobe(&probe_invoke);
	if (ret < 0) {
		pr_err("FAIL: register_kretprobe failed with %d\n", ret);
	} else {
		start_rec_time = ktime_get_coarse();
		pr_err("SMC_DBG: register_kretprobe OK\n");
	}

	return ret;
}

void unregister_symbol(void)
{
	struct smc_record *tmp, *traversal;
	unsigned long flags;

	unregister_kretprobe(&probe_invoke);

	pr_crit("Free nodes of smc_list\n");
	spin_lock_irqsave(&invoke_list_lock, flags);
	list_for_each_entry_safe(traversal, tmp, &invoke_record_head, node) {
			// remove inode record
			list_del(&traversal->node);
			kfree(traversal);
	}
	spin_unlock_irqrestore(&invoke_list_lock, flags);
	pr_crit("Free nodes of smc_list done\n");
}
