/* New client tool, replaces old /dev/initctl API and telinit tool
 *
 * Copyright (c) 2019-2025  Joachim Wiberg <troglobit@gmail.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

#include "config.h"

#include <dirent.h>
#include <stdio.h>
#include <signal.h>
#include <stdlib.h>
#include <search.h>
#include <inttypes.h>
#include <sys/sysinfo.h>		/* sysinfo() */

#include "cgutil.h"
#include "initctl.h"
#include "log.h"

#define CDIM   plain ? "" : "\e[2m"
#define CRST   plain ? "" : "\e[0m"
#define CLREOL plain ? "" : "\e[K"	/* Clear to end of line */

#define NONE " "
#define PIPE plain ? "| " : "│ "
#define FORK plain ? "|-" : "├─"
#define END  plain ? "`-" : "└─"

uint64_t total_ram;			/* From sysinfo() */

struct cg  dummy;			/* empty result "NULL"      */
struct cg *list;

int cgroup_avail(void)
{
	return fismnt(FINIT_CGPATH);
}

static size_t flen(FILE *fp)
{
	size_t total = 0, sz;
	char buf[512];

	while ((sz = fread(buf, 1, sizeof(buf), fp)) > 0)
		total += sz;
	rewind(fp);

	return total;
}

char *pid_cmdline(int pid)
{
	size_t i, len;
	char *buf;
	FILE *fp;

	fp = fopenf("r", "/proc/%d/cmdline", pid);
	if (!fp)
		return strdup("");	/* regular process */

	len = flen(fp);
	if (len == 0) {
	fail:
		fclose(fp);
		return NULL;		/* kernel thread */
	}

	buf = calloc(len + 1, 1);
	if (!buf)
		goto fail;

	if (fread(buf, 1, len, fp) < len) {
		free(buf);
		goto fail;
	}

	/* replace all NUL chars with space */
	for (i = 0; i < len; i++) {
		if (buf[i] == '\0')
			buf[i] = ' ';
	}

	fclose(fp);
	return buf;
}

char *pid_cgroup(int pid)
{
	char *buf, *ptr = NULL;
	size_t len;
	FILE *fp;

	fp = fopenf("r", "/proc/%d/cgroup", pid);
	if (!fp)
		return NULL;

	len = flen(fp);
	if (len == 0) {
		fclose(fp);
		return NULL;
	}
	len++;

	buf = calloc(1, len);
	if (!buf) {
		fclose(fp);
		return NULL;
	}

	if (fgets(buf, len, fp))
		ptr = chomp(buf);
	fclose(fp);

	if (ptr) {
		ptr = strchr(ptr, '/');
		if (ptr) {
			memmove(buf, ptr, strlen(ptr) + 1);
			return buf;
		}
	}

	free(buf);
	return NULL;
}

char *cgroup_val(char *path, char *file, char *buf, size_t len)
{
	char *val = NULL;
	FILE *fp;

	fp = fopenf("r", "%s/%s", path, file);
	if (fp) {
		if (fgets(buf, len, fp)) {
			val = chomp(buf);
			len = strcspn(val, " \t");
			val[len] = 0;
		}

		fclose(fp);
	}

	return val;
}

static uint64_t cgroup_uint64(char *path, char *file)
{
	uint64_t val = 0;
	char buf[42];

	if (cgroup_val(path, file, buf, sizeof(buf)))
		val = strtoull(buf, NULL, 10);

	return val;
}

static char *cgroup_memval(char *path, char *file, char *buf, size_t len)
{
	char data[42];

	if (cgroup_val(path, file, data, sizeof(data))) {
		if (!strcmp(data, "max"))
			strlcpy(buf, data, len);
		else
			memsz(strtoull(data, NULL, 10), buf, len);
	} else
		buf[0] = 0;

	return buf;
}

static uint64_t cgroup_memuse(struct cg *cg)
{
	char buf[42];
	FILE *fp;

	fp = fopenf("r", "%s/memory.stat", cg->cg_path);
	if (fp) {
		cg->cg_rss = 0;
		cg->cg_vmlib = 0;

		while (fgets(buf, sizeof(buf), fp)) {
			chomp(buf);

			if (!strncmp(buf, "anon", 4)) {
				cg->cg_rss += strtoull(&buf[5], NULL, 10);
				continue;
			}
			if (!strncmp(buf, "slab", 4)) {
				cg->cg_rss += strtoull(&buf[5], NULL, 10);
				continue;
			}
			if (!strncmp(buf, "kernel_stack", 12)) {
				cg->cg_rss += strtoull(&buf[5], NULL, 10);
				continue;
			}
			if (!strncmp(buf, "pagetables", 10)) {
				cg->cg_rss += strtoull(&buf[5], NULL, 10);
				continue;
			}
			if (!strncmp(buf, "percpu", 6)) {
				cg->cg_rss += strtoull(&buf[5], NULL, 10);
				continue;
			}
			if (!strncmp(buf, "sock", 4)) {
				cg->cg_rss += strtoull(&buf[5], NULL, 10);
				continue;
			}
			if (!strncmp(buf, "file", 4)) {
				cg->cg_vmlib += strtoull(&buf[5], NULL, 10);
				continue;
			}
			if (!strncmp(buf, "file_mapped", 11)) {
				cg->cg_vmlib += strtoull(&buf[5], NULL, 10);
				continue;
			}
		}
		fclose(fp);
	}

	cg->cg_memshare = (float)(cg->cg_rss * 100 / total_ram);

	return cg->cg_vmsize = cgroup_uint64(cg->cg_path, "memory.current");
}

uint64_t cgroup_memory(char *group)
{
	char path[256];

	paste(path, sizeof(path), FINIT_CGPATH, group);

	return cgroup_uint64(path, "memory.current");
}

int cgroup_throttle(char *group, uint64_t *throttled_usec, uint64_t *nr_throttled)
{
	char path[256];
	char buf[256];
	FILE *fp;
	int found = 0;

	paste(path, sizeof(path), FINIT_CGPATH, group);

	fp = fopenf("r", "%s/cpu.stat", path);
	if (!fp)
		return -1;

	*throttled_usec = 0;
	*nr_throttled = 0;

	while (fgets(buf, sizeof(buf), fp)) {
		chomp(buf);

		if (!strncmp(buf, "throttled_usec", 14)) {
			*throttled_usec = strtoull(&buf[15], NULL, 10);
			found++;
		} else if (!strncmp(buf, "nr_throttled", 12)) {
			*nr_throttled = strtoull(&buf[13], NULL, 10);
			found++;
		}

		if (found == 2)
			break;
	}

	fclose(fp);
	return 0;
}


static float cgroup_cpuload(struct cg *cg)
{
	char fn[256];
	char buf[64];
	FILE *fp;

	snprintf(fn, sizeof(fn), "%s/cpu.stat", cg->cg_path);
	fp = fopen(fn, "r");
	if (!fp)
		ERR(72, "Cannot open %s", fn);

	while (fgets(buf, sizeof(buf), fp)) {
		uint64_t curr;

		chomp(buf);
		if (strncmp(buf, "usage_usec", 10))
			continue;

		curr = strtoull(&buf[11], NULL, 10);
		if (cg->cg_prev != 0) {
			uint64_t diff = curr - cg->cg_prev;

			/* this expects 1 sec poll interval */
			cg->cg_load = (float)(diff / 1000000);
			cg->cg_load *= 100.0;
		}
		cg->cg_prev = curr;
		break;
	}

	fclose(fp);

	return cg->cg_load;
}

static struct cg *append(char *path)
{
	struct cg *cg;
	char fn[256];
	ENTRY item;

	snprintf(fn, sizeof(fn), "%s/cpu.stat", path);
	if (access(fn, F_OK)) {
		/* older kernels, 4.19, don't have summary cpu.stat in root */
		if (strcmp(path, FINIT_CGPATH))
			WARN("not a cgroup path with cpu controller, %s", path);
		return NULL;
	}

	cg = calloc(1, sizeof(struct cg));
	if (!cg)
		ERR(71, "failed allocating struct cg");

	cg->cg_path = strdup(path);
	if (list)
		cg->cg_next = list;
	list = cg;

	item.key  = cg->cg_path;
	item.data = cg;
	if (!hsearch(item, ENTER))
		ERR(70, "failed adding to hash table");

	return cg;
}

static struct cg *find(char *path)
{
	ENTRY *ep, item = { path, NULL };

	ep = hsearch(item, FIND);
	if (ep)
		return ep->data;

	return append(path);
}

/* update stats */
struct cg *cg_stats(char *path)
{
	struct cg *cg;

	cg = find(path);
	if (!cg)
		return &dummy;

	cgroup_cpuload(cg);
	cgroup_memuse(cg);

	return cg;
}

/*
 * Compare two memory limit values, returning the more restrictive one.
 * For memory.max: smaller value wins (except "max" means unlimited)
 * For memory.min: larger value wins (more protection)
 */
static uint64_t cmp_mem_limit(const char *a, const char *b, int is_max)
{
	uint64_t val_a, val_b;

	if (!a || !a[0])
		return b && b[0] ? strtoull(b, NULL, 10) : 0;
	if (!b || !b[0])
		return strtoull(a, NULL, 10);

	/* Handle "max" (unlimited) */
	if (!strcmp(a, "max"))
		return !strcmp(b, "max") ? UINT64_MAX : strtoull(b, NULL, 10);
	if (!strcmp(b, "max"))
		return strtoull(a, NULL, 10);

	val_a = strtoull(a, NULL, 10);
	val_b = strtoull(b, NULL, 10);

	if (is_max)
		return val_a < val_b ? val_a : val_b;  /* min for max limit */
	else
		return val_a > val_b ? val_a : val_b;  /* max for min limit */
}

/*
 * Compare two cpu.max values (format: "quota period")
 * Returns the more restrictive quota (smallest quota/period ratio)
 * Stores result in dst, up to len bytes
 */
static void cmp_cpu_max(const char *a, const char *b, char *dst, size_t len)
{
	char a_copy[32], b_copy[32];
	char *a_quota, *a_period, *b_quota, *b_period;
	uint64_t qa, pa, qb, pb;
	double ratio_a, ratio_b;

	if (!a || !a[0]) {
		if (b && b[0])
			strlcpy(dst, b, len);
		else
			dst[0] = 0;
		return;
	}
	if (!b || !b[0]) {
		strlcpy(dst, a, len);
		return;
	}

	/* Parse a */
	strlcpy(a_copy, a, sizeof(a_copy));
	a_quota = strtok(a_copy, " ");
	a_period = strtok(NULL, " ");

	/* Parse b */
	strlcpy(b_copy, b, sizeof(b_copy));
	b_quota = strtok(b_copy, " ");
	b_period = strtok(NULL, " ");

	/* Handle "max" quota (unlimited) */
	if (a_quota && !strcmp(a_quota, "max")) {
		strlcpy(dst, b, len);
		return;
	}
	if (b_quota && !strcmp(b_quota, "max")) {
		strlcpy(dst, a, len);
		return;
	}

	/* Compare ratios (quota/period) - smaller ratio is more restrictive */
	if (!a_quota || !a_period || !b_quota || !b_period) {
		strlcpy(dst, a, len);  /* fallback */
		return;
	}

	qa = strtoull(a_quota, NULL, 10);
	pa = strtoull(a_period, NULL, 10);
	qb = strtoull(b_quota, NULL, 10);
	pb = strtoull(b_period, NULL, 10);

	if (pa == 0 || pb == 0) {
		strlcpy(dst, a, len);  /* fallback */
		return;
	}

	ratio_a = (double)qa / (double)pa;
	ratio_b = (double)qb / (double)pb;

	strlcpy(dst, ratio_a <= ratio_b ? a : b, len);
}

/* query config with hierarchical limit resolution */
struct cg *cg_conf(char *path)
{
	static struct cg cg;
	char parent[512];
	char tmp[32];
	uint64_t mem_min_val, mem_max_val;

	/* Read initial values from the leaf cgroup */
	cgroup_val(path, "memory.min", cg.cg_mem.min, sizeof(cg.cg_mem.min));
	cgroup_val(path, "memory.max", cg.cg_mem.max, sizeof(cg.cg_mem.max));
	cgroup_val(path, "cpu.weight", cg.cg_cpu.weight, sizeof(cg.cg_cpu.weight));
	cgroup_val(path, "cpu.max",    cg.cg_cpu.max, sizeof(cg.cg_cpu.max));
	cgroup_val(path, "cpuset.cpus.effective", cg.cg_cpu.set, sizeof(cg.cg_cpu.set));
	cg.cg_vmsize = cgroup_uint64(path, "memory.current");

	/* Walk up the hierarchy to find most restrictive limits */
	strlcpy(parent, path, sizeof(parent));
	while (strcmp(parent, FINIT_CGPATH) != 0) {
		char *slash = strrchr(parent, '/');
		if (!slash || slash == parent)
			break;
		*slash = '\0';  /* Move up one level */

		/* Compare and update memory.min (take maximum) */
		if (cgroup_val(parent, "memory.min", tmp, sizeof(tmp))) {
			mem_min_val = cmp_mem_limit(cg.cg_mem.min, tmp, 0);
			if (mem_min_val == UINT64_MAX)
				strlcpy(cg.cg_mem.min, "max", sizeof(cg.cg_mem.min));
			else
				snprintf(cg.cg_mem.min, sizeof(cg.cg_mem.min), "%lu", mem_min_val);
		}

		/* Compare and update memory.max (take minimum) */
		if (cgroup_val(parent, "memory.max", tmp, sizeof(tmp))) {
			mem_max_val = cmp_mem_limit(cg.cg_mem.max, tmp, 1);
			if (mem_max_val == UINT64_MAX)
				strlcpy(cg.cg_mem.max, "max", sizeof(cg.cg_mem.max));
			else
				snprintf(cg.cg_mem.max, sizeof(cg.cg_mem.max), "%lu", mem_max_val);
		}

		/* Compare and update cpu.max (take most restrictive) */
		if (cgroup_val(parent, "cpu.max", tmp, sizeof(tmp)))
			cmp_cpu_max(cg.cg_cpu.max, tmp, cg.cg_cpu.max, sizeof(cg.cg_cpu.max));
	}

	return &cg;
}

static int cgroup_filter(const struct dirent *entry)
{
	/* Skip current dir ".", and prev dir "..", from list of files */
	if ((1 == strlen(entry->d_name) && entry->d_name[0] == '.') ||
	    (2 == strlen(entry->d_name) && !strcmp(entry->d_name, "..")))
		return 0;

	if (entry->d_name[0] == '.')
		return 0;

	if (entry->d_type != DT_DIR)
		return 0;

	return 1;
}

int cgroup_tree(char *path, char *pfx, int mode, int pos)
{
	struct dirent **namelist = NULL;
	char s[32], r[32], l[32], row[1024];
	size_t rlen = sizeof(row) - 1;
	struct stat st;
	struct cg *cg;
	char buf[512];
	FILE *fp;
	int i, n;
	int num;

	if (pos >= ttrows - 1)
		return pos;

	if (-1 == lstat(path, &st))
		return pos;

	if ((st.st_mode & S_IFMT) != S_IFDIR) {
		errno = ENOTDIR;
		return pos;
	}

	fp = fopenf("r", "%s/cgroup.procs", path);
	if (!fp)
		return pos;
	num = 0;
	while (fgets(buf, sizeof(buf), fp))
		num++;

	if (!pfx) {
		pfx = "";
		switch (mode) {
		case 1:
			cg = cg_stats(path);
			snprintf(row, rlen, " %6.6s  %6.6s  %6.6s %5.1f %5.1f  %s",
				 memsz(cg->cg_vmsize, s, sizeof(s)),
				 memsz(cg->cg_rss,    r, sizeof(r)),
				 memsz(cg->cg_vmlib,  l, sizeof(l)),
				 cg->cg_memshare, cg->cg_load, path);
			break;
		case 2:
			cg = cg_conf(path);
			snprintf(row, rlen, "%6.6s [%-6.6s%6.6s] %6s [%-6.6s%6.6s] %s",
				 memsz(cg->cg_vmsize, s, sizeof(s)),
				 cg->cg_mem.min, cg->cg_mem.max, cg->cg_cpu.set,
				 cg->cg_cpu.weight, cg->cg_cpu.max, path);
			break;
		default:
			strlcpy(row, path, rlen);
			break;
		}

		printf("\r%s%s\n", row, CLREOL);
		pos++;
		if (pos >= ttrows - 1)
			goto out;
	}

	if (num > 0) {
		rewind(fp);

		i = 0;
		while (fgets(buf, sizeof(buf), fp)) {
			char *cmdline;
			pid_t pid;

			pid = atoi(chomp(buf));
			if (pid <= 0)
				continue;

			/* skip kernel threads for now (no cmdline) */
			cmdline = pid_cmdline(pid);
			if (cmdline) {
				char proc[512];  /* Build full command, truncate row later */

				switch (mode) {
				case 1:
					snprintf(row, rlen, "%37s", " ");
					break;
				case 2:
					snprintf(row, rlen, " --.-- [            ]        [            ] ");
					break;
				default:
					row[0] = '\0';
					break;
				}

				strlcat(row, pfx, rlen);
				strlcat(row, ++i == num ? END : FORK, rlen);

				snprintf(proc, sizeof(proc), " %d %s", pid, cmdline);
				strlcat(row, CDIM, rlen);
				strlcat(row, proc, rlen);
				strlcat(row, CRST, sizeof(row));

				printf("\r%s%s\n", row, CLREOL);
				pos++;

				free(cmdline);

				if (pos >= ttrows - 1)
					break;
			}
		}
	}

out:
	fclose(fp);

	n = scandir(path, &namelist, cgroup_filter, alphasort);
	if (n > 0) {
		for (i = 0; i < n; i++) {
			char *nm = namelist[i]->d_name;
			char prefix[80];

			if (pos >= ttrows - 1)
				break;

			snprintf(buf, sizeof(buf), "%s/%s", path, nm);
			switch (mode) {
			case 1:
				cg = cg_stats(buf);
				snprintf(row, rlen,
					 " %6.6s  %6.6s  %6.6s %5.1f %5.1f  ",
					 memsz(cg->cg_vmsize, s, sizeof(s)),
					 memsz(cg->cg_rss,    r, sizeof(r)),
					 memsz(cg->cg_vmlib,  l, sizeof(l)),
					 cg->cg_memshare, cg->cg_load);
				break;
			case 2:
				cg = cg_conf(buf);
				snprintf(row, rlen, "%6.6s [%-6.6s%6.6s] %6.6s [%-6.6s%6.6s] ",
					 memsz(cg->cg_vmsize, s, sizeof(s)),
					 cg->cg_mem.min, cg->cg_mem.max,
					 cg->cg_cpu.set, cg->cg_cpu.weight, cg->cg_cpu.max);
				break;
			default:
				row[0] = '\0';
				break;
			}

			strlcat(row, pfx, rlen);
			if (i + 1 == n) {
				strlcat(row, END, rlen);
				snprintf(prefix, sizeof(prefix), "%s   ", pfx);
			} else {
				strlcat(row, FORK, rlen);
				snprintf(prefix, sizeof(prefix), "%s%s  ", pfx, PIPE);
			}
			strlcat(row, " ", rlen);

			strlcat(row, nm,   rlen);
			strlcat(row, "/ ", rlen);

			printf("\r%s%s\n", row, CLREOL);
			pos++;

			pos = cgroup_tree(buf, prefix, mode, pos);

			free(namelist[i]);
		}

		free(namelist);
	}

	return pos;
}

int show_cgps(char *arg)
{
	char path[512];

	if (!arg)
		arg = FINIT_CGPATH;
	else if (arg[0] != '/') {
		paste(path, sizeof(path), FINIT_CGPATH, arg);
		arg = path;
	}

	return cgroup_tree(arg, NULL, 0, 0);
}

static void cgtop(uev_t *w, void *arg, int events)
{
	static int first = 1;
	int lines;

	(void)w;
	(void)events;

	/* Re-probe screen size in case terminal was resized */
	ttinit(1);

	/* Clear screen on first run to remove any artifacts */
	if (first) {
		fputs("\e[?7l\e[2J\e[H", stdout);  /* Disable line wrap, clear, home */
		first = 0;
	} else {
		/* Move cursor to home instead of clearing screen to avoid flicker */
		fputs("\e[H", stdout);
	}

	lines = 0;
	if (heading) {
		print_header(" VmSIZE     RSS   VmLIB  %%MEM  %%CPU  GROUP");
		lines = 1;
	}
	cgroup_tree(arg, NULL, 1, lines);

	/* Clear from cursor to end of screen to remove leftover lines */
	fputs("\e[J", stdout);
	fflush(stdout);
}

static void cleanup(void)
{
	fputs("\e[?7h", stdout);	/* Re-enable line wrap */
	ttcooked();
	showcursor();
	puts("");
}

static void leave(uev_t *w, void *arg, int events)
{
	(void)arg;
	(void)events;

	uev_exit(w->ctx);
}

static void key(uev_t *w, void *arg, int events)
{
	char ch;

	(void)arg;
	(void)events;

	if (read(w->fd, &ch, sizeof(ch)) != -1) {
		switch (ch) {
		case 'q':
			uev_exit(w->ctx);
			break;

		default:
			dbg("Got char 0x%02x", ch);
			break;
		}
	}
}

int show_cgtop(char *arg)
{
	struct sysinfo si = { 0 };
        uev_t timer, input, sigint, sigterm, sigquit;
	char path[512];
        uev_ctx_t ctx;

	if (!arg)
		arg = FINIT_CGPATH;
	else if (arg[0] != '/') {
		paste(path, sizeof(path), FINIT_CGPATH, arg);
		arg = path;
	}

	if (!hcreate(ttrows + 25))
		ERR(70, "failed creating hash table");

	/* Ensure we have correct terminal size before starting */
	ttinit(1);

	sysinfo(&si);
	total_ram = si.totalram * si.mem_unit;

        uev_init(&ctx);
        uev_timer_init(&ctx, &timer, cgtop, arg, 1, ionce ? 0 : 1000);

	if (!ionce && !plain) {
		int flags;

		atexit(cleanup);
		ttraw();
		hidecursor();

		flags = fcntl(STDIN_FILENO, F_GETFL);
		if (flags != -1)
			(void)fcntl(STDIN_FILENO, F_SETFL, flags | O_NONBLOCK);
		(void)uev_io_init(&ctx, &input, key, NULL, STDIN_FILENO, UEV_READ);

		(void)uev_signal_init(&ctx, &sigint, leave, NULL, SIGINT);
		(void)uev_signal_init(&ctx, &sigterm, leave, NULL, SIGTERM);
		(void)uev_signal_init(&ctx, &sigquit, leave, NULL, SIGQUIT);
	}

	return uev_run(&ctx, 0);
}

int show_cgroup(char *arg)
{
	char path[512];

	if (!arg)
		arg = FINIT_CGPATH;
	else if (arg[0] != '/') {
		paste(path, sizeof(path), FINIT_CGPATH, arg);
		arg = path;
	}

	/* memory.current memory.min memory.max cpuset.cpus cpu.weight cpu.max */
	if (heading)
		print_header("   MEM [MIN      MAX]    CPU [WEIGHT   MAX] GROUP");

	return cgroup_tree(arg, NULL, 2, 0);
}

/**
 * Local Variables:
 *  indent-tabs-mode: t
 *  c-file-style: "linux"
 * End:
 */