// SPDX-License-Identifier: GPL-2.0

#include <linux/errno.h>
#include <linux/kernel.h>
#include <linux/mm.h>
#include <linux/module.h>
#include <linux/moduleparam.h>
#include <linux/sched.h>
#include <linux/slab.h>
#include <linux/string.h>
#include <linux/uaccess.h>
#include <linux/ktime.h>
#include <linux/xarray.h>
#include <linux/percpu_counter.h>
#include <linux/pid_namespace.h>
#include <linux/seq_file.h>
#include <linux/rcupdate.h>

#include <linux/sunrpc/clnt.h>
#include <linux/sunrpc/xdr.h>
#include <trace/events/sunrpc.h>

#include "jobid.h"


#define VAR_SEPARATOR "#"
#define MAX_VAR_COUNT 1
#define OUTPUT_PREFIX "#VastMD="

struct env_cache *jobid_cache;

/* Global cache structure */
struct env_cache {
	struct xarray xa;
	unsigned int max_entries;
	atomic_t cur_entries;
	struct percpu_counter hits;
	struct percpu_counter misses;
};

struct job_id_vars {
	int count;
	char *names[MAX_VAR_COUNT];
};

static struct job_id_vars __rcu *job_id_vars = NULL;


static inline bool is_module_unloading(void)
{
	if (!try_module_get(THIS_MODULE))
		return true;

	module_put(THIS_MODULE);
	return false;
}

static void synchronize_free_job_id_vars(struct job_id_vars *vars)
{
	int i;

	if (!vars)
		return;

	synchronize_rcu();
	for (i = 0; i < vars->count; i++)
		kfree(vars->names[i]);
	kfree(vars);
}

static int job_id_var_set(const char *val, const struct kernel_param *kp)
{
	struct job_id_vars *new_vars, *old_vars = NULL;
	char *input_str = NULL, *input_str_head = NULL, *token;
	int i = 0, ret = 0;

	if (is_module_unloading())
		return -ENODEV;

	new_vars = kzalloc(sizeof(*new_vars), GFP_KERNEL);
	if (!new_vars)
		return -ENOMEM;

	if (val && *val) {
		input_str = kstrdup(val, GFP_KERNEL);
		if (!input_str) {
			ret = -ENOMEM;
			goto out_free_new;
		}
		input_str_head = input_str;

		while ((token = strsep(&input_str, ","))) {
			token = strim(token);
			if (!*token)
				continue;
			if (i >= MAX_VAR_COUNT) {
				ret = -E2BIG;
				goto out_free_names;
			}
			new_vars->names[i] = kstrdup(token, GFP_KERNEL);
			if (!new_vars->names[i]) {
				ret = -ENOMEM;
				goto out_free_names;
			}
			i++;
		}
		new_vars->count = i;
		kfree(input_str_head);
	} else {
		new_vars->count = 0;
	}

	old_vars = rcu_dereference_protected(job_id_vars,
					     lockdep_is_held(&THIS_MODULE->param_lock));
	rcu_assign_pointer(job_id_vars, new_vars);

	synchronize_free_job_id_vars(old_vars);

	return 0;

out_free_names:
	for (int j = 0; j < i; j++)
		kfree(new_vars->names[j]);
	kfree(input_str_head);
out_free_new:
	kfree(new_vars);
	return ret;
}

static int job_id_var_get(char *buffer, const struct kernel_param *kp)
{
	struct job_id_vars *vars;
	size_t len = 0;

	rcu_read_lock();
	vars = rcu_dereference(job_id_vars);
	if (!vars) {
		rcu_read_unlock();
		return 0;
	}

	for (int i = 0; i < vars->count; i++) {
		const char *fmt = i > 0 ? ",%s" : "%s";
		int size = PAGE_SIZE - len;
		int written = snprintf(buffer + len, size, fmt, vars->names[i]);
		if (written >= size || written < 0) {
			rcu_read_unlock();
			return -E2BIG;
		}
		len += written;
	}
	rcu_read_unlock();
	return len;
}

static const struct kernel_param_ops job_id_var_ops = {
	.set = job_id_var_set,
	.get = job_id_var_get,
};

module_param_cb(job_id_var_names, &job_id_var_ops, NULL, 0644);
MODULE_PARM_DESC(job_id_var_names, "List of job ID environment names to scan for (comma-separated)");

static int jobid_cache_max_entries_set(const char *val, const struct kernel_param *kp)
{
	unsigned int new_val;
	int ret;

	if (is_module_unloading())
		return -ENODEV;

	ret = kstrtouint(val, 0, &new_val);
	if (ret)
		return ret;
	*((unsigned int *)kp->arg) = new_val;
	if (jobid_cache) {
		jobid_cache->max_entries = new_val;
		env_cache_flush();
	}
	return 0;
}

static int jobid_cache_max_entries_get(char *buffer, const struct kernel_param *kp)
{
	return sprintf(buffer, "%u\n", *((unsigned int *)kp->arg));
}

static const struct kernel_param_ops jobid_cache_max_entries_ops = {
	.set = jobid_cache_max_entries_set,
	.get = jobid_cache_max_entries_get,
};

static unsigned int jobid_cache_max_entries = 1024;
module_param_cb(jobid_cache_max_entries, &jobid_cache_max_entries_ops, &jobid_cache_max_entries, 0644);
MODULE_PARM_DESC(jobid_cache_max_entries, "Maximum number of entries in jobid cache. Cache flushed on write");

/* Cache entry for PID-to-environment mapping */
struct env_cache_entry {
	struct rcu_head rcu_head;
	struct list_head dispose;
	pid_t global_tgid;
	u64 start_time;
	atomic64_t last_access;
	char env_data[JOB_ID_MAX_LENGTH];
};

static ssize_t get_current_environ(char **out_env)
{
	struct mm_struct *mm;
	unsigned long env_start, env_end;
	char *buf = NULL;
	size_t total, copied = 0;
	ssize_t ret = 0;

	mm = get_task_mm(current);
	if (!mm)
		return -ESRCH;

	spin_lock(&mm->arg_lock);
	env_start = mm->env_start;
	env_end = mm->env_end;
	spin_unlock(&mm->arg_lock);

	if (!env_end) {
		mmput(mm);
		*out_env = NULL;
		return 0;
	}

	total = env_end - env_start;
	buf = kvmalloc(total + 1, GFP_KERNEL);
	if (!buf) {
		mmput(mm);
		return -ENOMEM;
	}

	while (copied < total) {
		size_t this_len = min_t(size_t, total - copied, PAGE_SIZE);
		unsigned long remaining;

		remaining = copy_from_user(buf + copied, 
		                           (void __user *)(env_start + copied), 
		                           this_len);
		
		if (remaining) {
			pr_warn_ratelimited("Failed capturing environment for Job-ID\n");
			ret = -EFAULT;
			goto out_free;
		}
		copied += this_len;
	}

	buf[total] = '\0';
	*out_env = buf;
	ret = total;

out_free:
	mmput(mm);
	if (ret < 0) {
		kvfree(buf);
		*out_env = NULL;
	}

	return ret;
}

static int env_var_valid(const char *env_var)
{
	struct job_id_vars *vars;
	size_t len;
	int i;
	int ret = -1;

	rcu_read_lock();
	vars = rcu_dereference(job_id_vars);
	if (!vars) {
		ret = -1;
		goto out;
	}

	for (i = 0; i < vars->count; i++) {
		if (!vars->names[i])
			continue;

		len = strlen(vars->names[i]);

		if (strncmp(env_var, vars->names[i], len) != 0)
			continue;
		if (env_var[len] != '=')
			continue;
		if (strstr(env_var, VAR_SEPARATOR))
			continue;

		ret = i;
		break;
	}

out:
	rcu_read_unlock();
	return ret;
}

static bool is_jobid_feature_enabled(void)
{
	struct job_id_vars *vars;
	bool ret;

	rcu_read_lock();
	vars = rcu_dereference(job_id_vars);
	if (!vars) {
		rcu_read_unlock();
		return false;
	}
	ret = vars->count > 0;
	rcu_read_unlock();
	return ret;
}

static void env_cache_entry_free(struct rcu_head *head)
{
	struct env_cache_entry *entry = container_of(head, struct env_cache_entry, rcu_head);
	kfree(entry);
}

static bool env_cache_lookup(pid_t global_tgid, u64 start_time, char *buf, size_t size)
{
	struct env_cache_entry *entry;
	bool found = false;
	u64 lookup_start = ktime_get_ns();

	if (!jobid_cache)
		return false;

	rcu_read_lock();
	entry = xa_load(&jobid_cache->xa, global_tgid);

	if (entry && entry->start_time == start_time) {
		strscpy(buf, entry->env_data, size);
		atomic64_set(&entry->last_access, lookup_start);
		percpu_counter_inc(&jobid_cache->hits);
		found = true;
	} else {
		percpu_counter_inc(&jobid_cache->misses);
	}
	rcu_read_unlock();

	trace_sunrpc_jobid_cache_lookup(global_tgid,
					start_time,
					found,
					ktime_get_ns() - lookup_start);

	return found;
}

/*
 * Evict the oldest entry from the cache based on last_access timestamp.
 * Must be called with xa_lock held.
 * Returns pointer to evicted entry for RCU-safe freeing by caller, or NULL.
 */
static struct env_cache_entry *env_cache_evict_oldest(void)
{
	unsigned long index;
	struct env_cache_entry *entry, *oldest = NULL;
	u64 now = ktime_get_ns();

	xa_for_each(&jobid_cache->xa, index, entry) {
		if (!oldest) {
			oldest = entry;
			continue;
		}
		u64 entry_access = atomic64_read(&entry->last_access);
		u64 oldest_access = atomic64_read(&oldest->last_access);
		if (entry_access < oldest_access)
			oldest = entry;
	}

	if (!oldest)
		return NULL;

	u64 age_ns = now - atomic64_read(&oldest->last_access);

	oldest = __xa_erase(&jobid_cache->xa, oldest->global_tgid);
	atomic_dec(&jobid_cache->cur_entries);
	trace_sunrpc_jobid_cache_evict(oldest->global_tgid, age_ns);

	return oldest;
}

static int env_cache_insert(pid_t global_tgid, u64 start_time, const char *env_data)
{
	struct env_cache_entry *new_entry, *old_entry, *evicted = NULL;
	int ret;
	unsigned int cur_entries;

	if (!jobid_cache)
		return -EINVAL;

	xa_lock(&jobid_cache->xa);

	if (atomic_read(&jobid_cache->cur_entries) >= jobid_cache->max_entries) {
		evicted = env_cache_evict_oldest();
	}

	new_entry = kzalloc(sizeof(*new_entry), GFP_ATOMIC);
	if (!new_entry) {
		ret = -ENOMEM;
		goto out;
	}

	new_entry->global_tgid = global_tgid;
	new_entry->start_time = start_time;
	atomic64_set(&new_entry->last_access, ktime_get_ns());
	INIT_LIST_HEAD(&new_entry->dispose);
	strscpy(new_entry->env_data, env_data, JOB_ID_MAX_LENGTH);

	/* Insert or replace */
	old_entry = __xa_store(&jobid_cache->xa, global_tgid, new_entry, GFP_ATOMIC);
	
	if (xa_is_err(old_entry)) {
		kfree(new_entry);
		ret = xa_err(old_entry);
		goto out;
	}

	if (!old_entry)
		/* This is a new entry - increment counter */
		atomic_inc(&jobid_cache->cur_entries);

	ret = 0;

out:
	cur_entries = atomic_read(&jobid_cache->cur_entries);
	xa_unlock(&jobid_cache->xa);

	if (ret == 0 && old_entry)
		call_rcu(&old_entry->rcu_head, env_cache_entry_free);

	if (evicted)
		call_rcu(&evicted->rcu_head, env_cache_entry_free);

	trace_sunrpc_jobid_cache_insert(global_tgid, ret == 0, cur_entries);
	return ret;
}

void env_cache_flush(void)
{
	LIST_HEAD(entries_to_dispose);
	struct env_cache_entry *entry, *tmp;
	unsigned long index;

	if (!jobid_cache)
		return;

	xa_lock(&jobid_cache->xa);
	xa_for_each(&jobid_cache->xa, index, entry) {
		list_add(&entry->dispose, &entries_to_dispose);
		__xa_erase(&jobid_cache->xa, index);
	}
	atomic_set(&jobid_cache->cur_entries, 0);
	percpu_counter_set(&jobid_cache->hits, 0);
	percpu_counter_set(&jobid_cache->misses, 0);
	xa_unlock(&jobid_cache->xa);

	synchronize_rcu();

	list_for_each_entry_safe(entry, tmp, &entries_to_dispose, dispose) {
		kfree(entry);
	}
}

void rpc_capture_task_env(char *buf, size_t size)
{
	char *env;
	char *env_var_ptr;
	char *valid_vars[MAX_VAR_COUNT] = {};
	bool buffer_valid = false;
	ssize_t env_len;
	size_t len = 0;
	int i;
	u64 start_time;
	pid_t global_tgid;
	u64 process_start_time;

	if (likely(!is_jobid_feature_enabled())) {
		buf[0] = '\0';
		return;
	}

	if (WARN_ON_ONCE(size < strlen(OUTPUT_PREFIX) + 1)) {
		buf[0] = '\0';
		return;
	}

	start_time = ktime_get_ns();

	global_tgid = task_tgid_nr_ns(current, &init_pid_ns);
	process_start_time = current->start_boottime;

	if (env_cache_lookup(global_tgid, process_start_time, buf, size))
		return;

	env_len = get_current_environ(&env);
	if (env_len <= 0 || !env) {
		buf[0] = '\0';
		return;
	}

	env_var_ptr = env;
	while (env_var_ptr < env + env_len && *env_var_ptr) {
		size_t var_len = strlen(env_var_ptr);
		int idx;

		idx = env_var_valid(env_var_ptr);
		if (idx >= 0)
			valid_vars[idx] = env_var_ptr;

		env_var_ptr += var_len + 1;
	}

	len = snprintf(buf, size, OUTPUT_PREFIX);
	if (len < 0) {
		buffer_valid = false;
		goto out;
	}

	for (i = 0; i < MAX_VAR_COUNT; i++) {
		if (valid_vars[i]) {
			char *eq_char = strchr(valid_vars[i], '=');
			char *env_var_value;
			int needed_space;
			size_t available_space = size - len;

			if (!eq_char)
				continue;
			env_var_value = eq_char + 1;

			needed_space = snprintf(buf + len, available_space, "%s#", env_var_value);
			if (needed_space < 0) {
				buffer_valid = false;
				break;
			}
			if (needed_space < available_space) {
				len += needed_space;
				buffer_valid = true;
			}
		}
	}

out:
	if (!buffer_valid){
		buf[0] = '\0';
	} else {
		buf[len] = '\0';
		trace_sunrpc_jobid_capture(buf, current->pid,
					   ktime_get_ns() - start_time);
		env_cache_insert(global_tgid, process_start_time, buf);
	}

	kvfree(env);
}
EXPORT_SYMBOL_GPL(rpc_capture_task_env);

char *alloc_jobid_buffer(void)
{
	if (likely(!is_jobid_feature_enabled()))
		return NULL;

	return kmalloc(JOB_ID_MAX_LENGTH, GFP_KERNEL);
}

int marshal_nodename_and_job_id(struct rpc_task *task, struct xdr_stream *xdr)
{
	struct rpc_clnt *clnt = task->tk_client;
	char new_name[RPC_MAX_MACHINENAME];
	int job_len = strlen(task->tk_job_id);
	int total_len = clnt->cl_nodelen + job_len;

	if (total_len > sizeof(new_name))
		return -EMSGSIZE;

	memcpy(new_name, clnt->cl_nodename, clnt->cl_nodelen);
	memcpy(new_name + clnt->cl_nodelen, task->tk_job_id, job_len);

	return xdr_stream_encode_opaque(xdr, new_name, total_len);
}

int jobid_init(void)
{
	int err;

	jobid_cache = kzalloc(sizeof(*jobid_cache), GFP_KERNEL);
	if (!jobid_cache) {
		err = -ENOMEM;
		goto out_clean_job_id_vars;
	}

	xa_init(&jobid_cache->xa);
	jobid_cache->max_entries = jobid_cache_max_entries;
	atomic_set(&jobid_cache->cur_entries, 0);

	err = percpu_counter_init(&jobid_cache->hits, 0, GFP_KERNEL);
	if (err)
		goto out_free_cache;

	err = percpu_counter_init(&jobid_cache->misses, 0, GFP_KERNEL);
	if (err)
		goto out_destroy_hits;

	return 0;

out_destroy_hits:
	percpu_counter_destroy(&jobid_cache->hits);
out_free_cache:
	xa_destroy(&jobid_cache->xa);
	kfree(jobid_cache);
	jobid_cache = NULL;
out_clean_job_id_vars:
	/* Free Jobid vars that might have been set via module parameter */
	synchronize_free_job_id_vars(rcu_access_pointer(job_id_vars));
	rcu_assign_pointer(job_id_vars, NULL);
	return err;
}

void jobid_destroy(void)
{
	struct job_id_vars *vars;

	/* Taking param lock to prevent races with parameter setters*/
	kernel_param_lock(THIS_MODULE);

	vars = rcu_dereference_protected(job_id_vars,
					 lockdep_is_held(&THIS_MODULE->param_lock));
	rcu_assign_pointer(job_id_vars, NULL);
	synchronize_free_job_id_vars(vars);

	if (!jobid_cache)
		goto out;

	env_cache_flush();
	xa_destroy(&jobid_cache->xa);
	percpu_counter_destroy(&jobid_cache->hits);
	percpu_counter_destroy(&jobid_cache->misses);
	kfree(jobid_cache);
	jobid_cache = NULL;

out:
	kernel_param_unlock(THIS_MODULE);
}

void jobid_cache_print_stats(struct seq_file *m)
{
	s64 hits, misses, total;
	unsigned int cur_entries, max_entries;

	if (!jobid_cache) {
		seq_puts(m, "Cache not initialized\n");
		return;
	}

	hits = percpu_counter_sum(&jobid_cache->hits);
	misses = percpu_counter_sum(&jobid_cache->misses);
	total = hits + misses;
	cur_entries = atomic_read(&jobid_cache->cur_entries);
	max_entries = jobid_cache->max_entries;

	seq_printf(m, "Hits: %lld\n", hits);
	seq_printf(m, "Misses: %lld\n", misses);
	seq_printf(m, "Total lookups: %lld\n", total);
	if (total > 0)
		seq_printf(m, "Hit ratio: %lld.%02lld%%\n",
			   (hits * 100) / total,
			   ((hits * 10000) / total) % 100);
	else
		seq_puts(m, "Hit ratio: N/A\n");
	seq_printf(m, "Current entries: %u\n", cur_entries);
	seq_printf(m, "Max entries: %u\n", max_entries);
}

void jobid_cache_print_dump(struct seq_file *m)
{
	struct env_cache_entry *entry;
	unsigned long index;
	u64 now;

	if (!jobid_cache) {
		seq_puts(m, "Cache not initialized\n");
		return;
	}

	now = ktime_get_ns();
	seq_printf(m, "%-10s %-20s %-20s %-20s %s\n",
		   "TGID", "Start Time", "Last Access", "Age (ms)", "Env Data");

	rcu_read_lock();
	xa_for_each(&jobid_cache->xa, index, entry) {
		u64 last_access = atomic64_read(&entry->last_access);
		u64 age_ns = now - last_access;
		u64 age_ms = age_ns / 1000000;

		seq_printf(m, "%-10d %-20llu %-20llu %-20llu %s\n",
			   entry->global_tgid,
			   entry->start_time,
			   last_access,
			   age_ms,
			   entry->env_data);
	}
	rcu_read_unlock();
}
