// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2025, Google LLC.
 */

#include <time.h>

#include "lru_gen_util.h"

/*
 * Tracks state while we parse memcg lru_gen stats. The file we're parsing is
 * structured like this (some extra whitespace elided):
 *
 * memcg (id) (path)
 * node (id)
 * (gen_nr) (age_in_ms) (nr_anon_pages) (nr_file_pages)
 */
struct memcg_stats_parse_context {
	bool consumed; /* Whether or not this line was consumed */
	/* Next parse handler to invoke */
	void (*next_handler)(struct memcg_stats *stats,
			     struct memcg_stats_parse_context *ctx,
			     char *line);
	int current_node_idx; /* Current index in nodes array */
	const char *name; /* The name of the memcg we're looking for */
};

static void memcg_stats_handle_searching(struct memcg_stats *stats,
					 struct memcg_stats_parse_context *ctx,
					 char *line);
static void memcg_stats_handle_in_memcg(struct memcg_stats *stats,
					struct memcg_stats_parse_context *ctx,
					char *line);
static void memcg_stats_handle_in_node(struct memcg_stats *stats,
				       struct memcg_stats_parse_context *ctx,
				       char *line);

struct split_iterator {
	char *str;
	char *save;
};

static char *split_next(struct split_iterator *it)
{
	char *ret = strtok_r(it->str, " \t\n\r", &it->save);

	it->str = NULL;
	return ret;
}

static void memcg_stats_handle_searching(struct memcg_stats *stats,
					 struct memcg_stats_parse_context *ctx,
					 char *line)
{
	struct split_iterator it = { .str = line };
	char *prefix = split_next(&it);
	char *memcg_id = split_next(&it);
	char *memcg_name = split_next(&it);
	char *end;

	ctx->consumed = true;

	if (!prefix || strcmp("memcg", prefix))
		return; /* Not a memcg line (maybe empty), skip */

	TEST_ASSERT(memcg_id && memcg_name,
		    "malformed memcg line; no memcg id or memcg_name");

	if (strcmp(memcg_name + 1, ctx->name))
		return; /* Wrong memcg, skip */

	/* Found it! */

	stats->memcg_id = strtoul(memcg_id, &end, 10);
	TEST_ASSERT(*end == '\0', "malformed memcg id '%s'", memcg_id);
	if (!stats->memcg_id)
		return; /* Removed memcg? */

	ctx->next_handler = memcg_stats_handle_in_memcg;
}

static void memcg_stats_handle_in_memcg(struct memcg_stats *stats,
					struct memcg_stats_parse_context *ctx,
					char *line)
{
	struct split_iterator it = { .str = line };
	char *prefix = split_next(&it);
	char *id = split_next(&it);
	long found_node_id;
	char *end;

	ctx->consumed = true;
	ctx->current_node_idx = -1;

	if (!prefix)
		return; /* Skip empty lines */

	if (!strcmp("memcg", prefix)) {
		/* Memcg done, found next one; stop. */
		ctx->next_handler = NULL;
		return;
	} else if (strcmp("node", prefix))
		TEST_ASSERT(false, "found malformed line after 'memcg ...',"
				   "token: '%s'", prefix);

	/* At this point we know we have a node line. Parse the ID. */

	TEST_ASSERT(id, "malformed node line; no node id");

	found_node_id = strtol(id, &end, 10);
	TEST_ASSERT(*end == '\0', "malformed node id '%s'", id);

	ctx->current_node_idx = stats->nr_nodes++;
	TEST_ASSERT(ctx->current_node_idx < MAX_NR_NODES,
		    "memcg has stats for too many nodes, max is %d",
		    MAX_NR_NODES);
	stats->nodes[ctx->current_node_idx].node = found_node_id;

	ctx->next_handler = memcg_stats_handle_in_node;
}

static void memcg_stats_handle_in_node(struct memcg_stats *stats,
				       struct memcg_stats_parse_context *ctx,
				       char *line)
{
	char *my_line = strdup(line);
	struct split_iterator it = { .str = my_line };
	char *gen, *age, *nr_anon, *nr_file;
	struct node_stats *node_stats;
	struct generation_stats *gen_stats;
	char *end;

	TEST_ASSERT(it.str, "failed to copy input line");

	gen = split_next(&it);

	if (!gen)
		goto out_consume; /* Skip empty lines */

	if (!strcmp("memcg", gen) || !strcmp("node", gen)) {
		/*
		 * Reached next memcg or node section. Don't consume, let the
		 * other handler deal with this.
		 */
		ctx->next_handler = memcg_stats_handle_in_memcg;
		goto out;
	}

	node_stats = &stats->nodes[ctx->current_node_idx];
	TEST_ASSERT(node_stats->nr_gens < MAX_NR_GENS,
		    "found too many generation lines; max is %d",
		    MAX_NR_GENS);
	gen_stats = &node_stats->gens[node_stats->nr_gens++];

	age = split_next(&it);
	nr_anon = split_next(&it);
	nr_file = split_next(&it);

	TEST_ASSERT(age && nr_anon && nr_file,
		    "malformed generation line; not enough tokens");

	gen_stats->gen = (int)strtol(gen, &end, 10);
	TEST_ASSERT(*end == '\0', "malformed generation number '%s'", gen);

	gen_stats->age_ms = strtol(age, &end, 10);
	TEST_ASSERT(*end == '\0', "malformed generation age '%s'", age);

	gen_stats->nr_anon = strtol(nr_anon, &end, 10);
	TEST_ASSERT(*end == '\0', "malformed anonymous page count '%s'",
		    nr_anon);

	gen_stats->nr_file = strtol(nr_file, &end, 10);
	TEST_ASSERT(*end == '\0', "malformed file page count '%s'", nr_file);

out_consume:
	ctx->consumed = true;
out:
	free(my_line);
}

static void print_memcg_stats(const struct memcg_stats *stats, const char *name)
{
	int node, gen;

	pr_debug("stats for memcg %s (id %lu):\n", name, stats->memcg_id);
	for (node = 0; node < stats->nr_nodes; ++node) {
		pr_debug("\tnode %d\n", stats->nodes[node].node);
		for (gen = 0; gen < stats->nodes[node].nr_gens; ++gen) {
			const struct generation_stats *gstats =
				&stats->nodes[node].gens[gen];

			pr_debug("\t\tgen %d\tage_ms %ld"
				 "\tnr_anon %ld\tnr_file %ld\n",
				 gstats->gen, gstats->age_ms, gstats->nr_anon,
				 gstats->nr_file);
		}
	}
}

/* Re-read lru_gen debugfs information for @memcg into @stats. */
void lru_gen_read_memcg_stats(struct memcg_stats *stats, const char *memcg)
{
	FILE *f;
	ssize_t read = 0;
	char *line = NULL;
	size_t bufsz;
	struct memcg_stats_parse_context ctx = {
		.next_handler = memcg_stats_handle_searching,
		.name = memcg,
	};

	memset(stats, 0, sizeof(struct memcg_stats));

	f = fopen(LRU_GEN_DEBUGFS, "r");
	TEST_ASSERT(f, "fopen(%s) failed", LRU_GEN_DEBUGFS);

	while (ctx.next_handler && (read = getline(&line, &bufsz, f)) > 0) {
		ctx.consumed = false;

		do {
			ctx.next_handler(stats, &ctx, line);
			if (!ctx.next_handler)
				break;
		} while (!ctx.consumed);
	}

	if (read < 0 && !feof(f))
		TEST_ASSERT(false, "getline(%s) failed", LRU_GEN_DEBUGFS);

	TEST_ASSERT(stats->memcg_id > 0, "Couldn't find memcg: %s\n"
		    "Did the memcg get created in the proper mount?",
		    memcg);
	if (line)
		free(line);
	TEST_ASSERT(!fclose(f), "fclose(%s) failed", LRU_GEN_DEBUGFS);

	print_memcg_stats(stats, memcg);
}

/*
 * Find all pages tracked by lru_gen for this memcg in generation @target_gen.
 *
 * If @target_gen is negative, look for all generations.
 */
long lru_gen_sum_memcg_stats_for_gen(int target_gen,
				     const struct memcg_stats *stats)
{
	int node, gen;
	long total_nr = 0;

	for (node = 0; node < stats->nr_nodes; ++node) {
		const struct node_stats *node_stats = &stats->nodes[node];

		for (gen = 0; gen < node_stats->nr_gens; ++gen) {
			const struct generation_stats *gen_stats =
				&node_stats->gens[gen];

			if (target_gen >= 0 && gen_stats->gen != target_gen)
				continue;

			total_nr += gen_stats->nr_anon + gen_stats->nr_file;
		}
	}

	return total_nr;
}

/* Find all pages tracked by lru_gen for this memcg. */
long lru_gen_sum_memcg_stats(const struct memcg_stats *stats)
{
	return lru_gen_sum_memcg_stats_for_gen(-1, stats);
}

/*
 * If lru_gen aging should force page table scanning.
 *
 * If you want to set this to false, you will need to do eviction
 * before doing extra aging passes.
 */
static const bool force_scan = true;

static void run_aging_impl(unsigned long memcg_id, int node_id, int max_gen)
{
	FILE *f = fopen(LRU_GEN_DEBUGFS, "w");
	char *command;
	size_t sz;

	TEST_ASSERT(f, "fopen(%s) failed", LRU_GEN_DEBUGFS);
	sz = asprintf(&command, "+ %lu %d %d 1 %d\n",
		      memcg_id, node_id, max_gen, force_scan);
	TEST_ASSERT(sz > 0, "creating aging command failed");

	pr_debug("Running aging command: %s", command);
	if (fwrite(command, sizeof(char), sz, f) < sz) {
		TEST_ASSERT(false, "writing aging command %s to %s failed",
			    command, LRU_GEN_DEBUGFS);
	}

	TEST_ASSERT(!fclose(f), "fclose(%s) failed", LRU_GEN_DEBUGFS);
}

void lru_gen_do_aging(struct memcg_stats *stats, const char *memcg)
{
	int node, gen;

	pr_debug("lru_gen: invoking aging...\n");

	/* Must read memcg stats to construct the proper aging command. */
	lru_gen_read_memcg_stats(stats, memcg);

	for (node = 0; node < stats->nr_nodes; ++node) {
		int max_gen = 0;

		for (gen = 0; gen < stats->nodes[node].nr_gens; ++gen) {
			int this_gen = stats->nodes[node].gens[gen].gen;

			max_gen = max_gen > this_gen ? max_gen : this_gen;
		}

		run_aging_impl(stats->memcg_id, stats->nodes[node].node,
			       max_gen);
	}

	/* Re-read so callers get updated information */
	lru_gen_read_memcg_stats(stats, memcg);
}

/*
 * Find which generation contains at least @pages pages, assuming that
 * such a generation exists.
 */
int lru_gen_find_generation(const struct memcg_stats *stats,
			    unsigned long pages)
{
	int node, gen, gen_idx, min_gen = INT_MAX, max_gen = -1;

	for (node = 0; node < stats->nr_nodes; ++node)
		for (gen_idx = 0; gen_idx < stats->nodes[node].nr_gens;
		     ++gen_idx) {
			gen = stats->nodes[node].gens[gen_idx].gen;
			max_gen = gen > max_gen ? gen : max_gen;
			min_gen = gen < min_gen ? gen : min_gen;
		}

	for (gen = min_gen; gen <= max_gen; ++gen)
		/* See if this generation has enough pages. */
		if (lru_gen_sum_memcg_stats_for_gen(gen, stats) > pages)
			return gen;

	return -1;
}

bool lru_gen_usable(void)
{
	long required_features = LRU_GEN_ENABLED | LRU_GEN_MM_WALK;
	int lru_gen_fd, lru_gen_debug_fd;
	char mglru_feature_str[8] = {};
	long mglru_features;

	lru_gen_fd = open(LRU_GEN_ENABLED_PATH, O_RDONLY);
	if (lru_gen_fd < 0) {
		puts("lru_gen: Could not open " LRU_GEN_ENABLED_PATH);
		return false;
	}
	if (read(lru_gen_fd, &mglru_feature_str, 7) < 7) {
		puts("lru_gen: Could not read from " LRU_GEN_ENABLED_PATH);
		close(lru_gen_fd);
		return false;
	}
	close(lru_gen_fd);

	mglru_features = strtol(mglru_feature_str, NULL, 16);
	if ((mglru_features & required_features) != required_features) {
		printf("lru_gen: missing features, got: 0x%lx, expected: 0x%lx\n",
		       mglru_features, required_features);
		printf("lru_gen: Try 'echo 0x%lx > /sys/kernel/mm/lru_gen/enabled'\n",
		       required_features);
		return false;
	}

	lru_gen_debug_fd = open(LRU_GEN_DEBUGFS, O_RDWR);
	__TEST_REQUIRE(lru_gen_debug_fd >= 0,
		       "lru_gen: Could not open " LRU_GEN_DEBUGFS ", "
		       "but lru_gen is enabled, so cannot use page_idle.");
	close(lru_gen_debug_fd);
	return true;
}