// SPDX-License-Identifier: GPL-2.0
/*
 * delaytop.c - system-wide delay monitoring tool.
 *
 * This tool provides real-time monitoring and statistics of
 * system, container, and task-level delays, including CPU,
 * memory, IO, and IRQ. It supports both interactive (top-like),
 * and can output delay information for the whole system, specific
 * containers (cgroups), or individual tasks (PIDs).
 *
 * Key features:
 *	- Collects per-task delay accounting statistics via taskstats.
 *	- Collects system-wide PSI information.
 *	- Supports sorting, filtering.
 *	- Supports both interactive (screen refresh).
 *
 * Copyright (C) Fan Yu, ZTE Corp. 2025
 * Copyright (C) Wang Yaxin, ZTE Corp. 2025
 *
 * Compile with
 *	gcc -I/usr/src/linux/include delaytop.c -o delaytop
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <getopt.h>
#include <signal.h>
#include <time.h>
#include <dirent.h>
#include <ctype.h>
#include <stdbool.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <termios.h>
#include <limits.h>
#include <linux/genetlink.h>
#include <linux/taskstats.h>
#include <linux/cgroupstats.h>
#include <stddef.h>

#define PSI_PATH	"/proc/pressure"
#define PSI_CPU_PATH	"/proc/pressure/cpu"
#define PSI_MEMORY_PATH	"/proc/pressure/memory"
#define PSI_IO_PATH	"/proc/pressure/io"
#define PSI_IRQ_PATH	"/proc/pressure/irq"

#define NLA_NEXT(na)			((struct nlattr *)((char *)(na) + NLA_ALIGN((na)->nla_len)))
#define NLA_DATA(na)			((void *)((char *)(na) + NLA_HDRLEN))
#define NLA_PAYLOAD(len)		(len - NLA_HDRLEN)

#define GENLMSG_DATA(glh)		((void *)(NLMSG_DATA(glh) + GENL_HDRLEN))
#define GENLMSG_PAYLOAD(glh)	(NLMSG_PAYLOAD(glh, 0) - GENL_HDRLEN)

#define TASK_COMM_LEN	16
#define MAX_MSG_SIZE	1024
#define MAX_TASKS		1000
#define MAX_BUF_LEN		256
#define SET_TASK_STAT(task_count, field) tasks[task_count].field = stats.field
#define BOOL_FPRINT(stream, fmt, ...) \
({ \
	int ret = fprintf(stream, fmt, ##__VA_ARGS__); \
	ret >= 0; \
})
#define TASK_AVG(task, field) average_ms((task).field##_delay_total, (task).field##_count)
#define PSI_LINE_FORMAT "%-12s %6.1f%%/%6.1f%%/%6.1f%%/%8llu(ms)\n"
#define DELAY_FMT_DEFAULT "%8.2f %8.2f %8.2f %8.2f\n"
#define DELAY_FMT_MEMVERBOSE "%8.2f %8.2f %8.2f %8.2f %8.2f %8.2f\n"
#define SORT_FIELD(name, cmd, modes) \
	{#name, #cmd, \
	offsetof(struct task_info, name##_delay_total), \
	offsetof(struct task_info, name##_count), \
	modes}
#define END_FIELD {NULL, 0, 0}

/* Display mode types */
#define MODE_TYPE_ALL	(0xFFFFFFFF)
#define MODE_DEFAULT	(1 << 0)
#define MODE_MEMVERBOSE	(1 << 1)

/* PSI statistics structure */
struct psi_stats {
	double cpu_some_avg10, cpu_some_avg60, cpu_some_avg300;
	unsigned long long cpu_some_total;
	double cpu_full_avg10, cpu_full_avg60, cpu_full_avg300;
	unsigned long long cpu_full_total;
	double memory_some_avg10, memory_some_avg60, memory_some_avg300;
	unsigned long long memory_some_total;
	double memory_full_avg10, memory_full_avg60, memory_full_avg300;
	unsigned long long memory_full_total;
	double io_some_avg10, io_some_avg60, io_some_avg300;
	unsigned long long io_some_total;
	double io_full_avg10, io_full_avg60, io_full_avg300;
	unsigned long long io_full_total;
	double irq_full_avg10, irq_full_avg60, irq_full_avg300;
	unsigned long long irq_full_total;
};

/* Task delay information structure */
struct task_info {
	int pid;
	int tgid;
	char command[TASK_COMM_LEN];
	unsigned long long cpu_count;
	unsigned long long cpu_delay_total;
	unsigned long long blkio_count;
	unsigned long long blkio_delay_total;
	unsigned long long swapin_count;
	unsigned long long swapin_delay_total;
	unsigned long long freepages_count;
	unsigned long long freepages_delay_total;
	unsigned long long thrashing_count;
	unsigned long long thrashing_delay_total;
	unsigned long long compact_count;
	unsigned long long compact_delay_total;
	unsigned long long wpcopy_count;
	unsigned long long wpcopy_delay_total;
	unsigned long long irq_count;
	unsigned long long irq_delay_total;
	unsigned long long mem_count;
	unsigned long long mem_delay_total;
};

/* Container statistics structure */
struct container_stats {
	int nr_sleeping;		/* Number of sleeping processes */
	int nr_running;			/* Number of running processes */
	int nr_stopped;			/* Number of stopped processes */
	int nr_uninterruptible; /* Number of uninterruptible processes */
	int nr_io_wait;			/* Number of processes in IO wait */
};

/* Delay field structure */
struct field_desc {
	const char *name;	/* Field name for cmdline argument */
	const char *cmd_char;	/* Interactive command */
	unsigned long total_offset; /* Offset of total delay in task_info */
	unsigned long count_offset; /* Offset of count in task_info */
	size_t supported_modes; /* Supported display modes */
};

/* Program settings structure */
struct config {
	int delay;				/* Update interval in seconds */
	int iterations;			/* Number of iterations, 0 == infinite */
	int max_processes;		/* Maximum number of processes to show */
	int output_one_time;	/* Output once and exit */
	int monitor_pid;		/* Monitor specific PID */
	char *container_path;	/* Path to container cgroup */
	const struct field_desc *sort_field;	/* Current sort field */
	size_t display_mode;	/* Current display mode */
};

/* Global variables */
static struct config cfg;
static struct psi_stats psi;
static struct task_info tasks[MAX_TASKS];
static int task_count;
static int running = 1;
static struct container_stats container_stats;
static const struct field_desc sort_fields[] = {
	SORT_FIELD(cpu,		c,	MODE_DEFAULT),
	SORT_FIELD(blkio,	i,	MODE_DEFAULT),
	SORT_FIELD(irq,		q,	MODE_DEFAULT),
	SORT_FIELD(mem,		m,	MODE_DEFAULT | MODE_MEMVERBOSE),
	SORT_FIELD(swapin,	s,	MODE_MEMVERBOSE),
	SORT_FIELD(freepages,	r,	MODE_MEMVERBOSE),
	SORT_FIELD(thrashing,	t,	MODE_MEMVERBOSE),
	SORT_FIELD(compact,	p,	MODE_MEMVERBOSE),
	SORT_FIELD(wpcopy,	w,	MODE_MEMVERBOSE),
	END_FIELD
};
static int sort_selected;

/* Netlink socket variables */
static int nl_sd = -1;
static int family_id;

/* Set terminal to non-canonical mode for q-to-quit */
static struct termios orig_termios;
static void enable_raw_mode(void)
{
	struct termios raw;

	tcgetattr(STDIN_FILENO, &orig_termios);
	raw = orig_termios;
	raw.c_lflag &= ~(ICANON | ECHO);
	tcsetattr(STDIN_FILENO, TCSAFLUSH, &raw);
}
static void disable_raw_mode(void)
{
	tcsetattr(STDIN_FILENO, TCSAFLUSH, &orig_termios);
}

/* Find field descriptor by command line */
static const struct field_desc *get_field_by_cmd_char(char ch)
{
	const struct field_desc *field;

	for (field = sort_fields; field->name != NULL; field++) {
		if (field->cmd_char[0] == ch)
			return field;
	}

	return NULL;
}

/* Find field descriptor by name with string comparison */
static const struct field_desc *get_field_by_name(const char *name)
{
	const struct field_desc *field;
	size_t field_len;

	for (field = sort_fields; field->name != NULL; field++) {
		field_len = strlen(field->name);
		if (field_len != strlen(name))
			continue;
		if (strncmp(field->name, name, field_len) == 0)
			return field;
	}

	return NULL;
}

/* Find display name for a field descriptor */
static const char *get_name_by_field(const struct field_desc *field)
{
	return field ? field->name : "UNKNOWN";
}

/* Generate string of available field names */
static void display_available_fields(size_t mode)
{
	const struct field_desc *field;
	char buf[MAX_BUF_LEN];

	buf[0] = '\0';

	for (field = sort_fields; field->name != NULL; field++) {
		if (!(field->supported_modes & mode))
			continue;
		strncat(buf, "|", MAX_BUF_LEN - strlen(buf) - 1);
		strncat(buf, field->name, MAX_BUF_LEN - strlen(buf) - 1);
		buf[MAX_BUF_LEN - 1] = '\0';
	}

	fprintf(stderr, "Available fields: %s\n", buf);
}

/* Display usage information and command line options */
static void usage(void)
{
	printf("Usage: delaytop [Options]\n"
	"Options:\n"
	"  -h, --help               Show this help message and exit\n"
	"  -d, --delay=SECONDS      Set refresh interval (default: 2 seconds, min: 1)\n"
	"  -n, --iterations=COUNT   Set number of updates (default: 0 = infinite)\n"
	"  -P, --processes=NUMBER   Set maximum number of processes to show (default: 20, max: 1000)\n"
	"  -o, --once               Display once and exit\n"
	"  -p, --pid=PID            Monitor only the specified PID\n"
	"  -C, --container=PATH     Monitor the container at specified cgroup path\n"
	"  -s, --sort=FIELD         Sort by delay field (default: cpu)\n"
	"  -M, --memverbose         Display memory detailed information\n");
	exit(0);
}

/* Parse command line arguments and set configuration */
static void parse_args(int argc, char **argv)
{
	int c;
	const struct field_desc *field;
	struct option long_options[] = {
		{"help", no_argument, 0, 'h'},
		{"delay", required_argument, 0, 'd'},
		{"iterations", required_argument, 0, 'n'},
		{"pid", required_argument, 0, 'p'},
		{"once", no_argument, 0, 'o'},
		{"processes", required_argument, 0, 'P'},
		{"sort", required_argument, 0, 's'},
		{"container", required_argument, 0, 'C'},
		{"memverbose", no_argument, 0, 'M'},
		{0, 0, 0, 0}
	};

	/* Set defaults */
	cfg.delay = 2;
	cfg.iterations = 0;
	cfg.max_processes = 20;
	cfg.sort_field = &sort_fields[0];	/* Default sorted by CPU delay */
	cfg.output_one_time = 0;
	cfg.monitor_pid = 0;	/* 0 means monitor all PIDs */
	cfg.container_path = NULL;
	cfg.display_mode = MODE_DEFAULT;

	while (1) {
		int option_index = 0;

		c = getopt_long(argc, argv, "hd:n:p:oP:C:s:M", long_options, &option_index);
		if (c == -1)
			break;

		switch (c) {
		case 'h':
			usage();
			break;
		case 'd':
			cfg.delay = atoi(optarg);
			if (cfg.delay < 1) {
				fprintf(stderr, "Error: delay must be >= 1.\n");
				exit(1);
			}
			break;
		case 'n':
			cfg.iterations = atoi(optarg);
			if (cfg.iterations < 0) {
				fprintf(stderr, "Error: iterations must be >= 0.\n");
				exit(1);
			}
			break;
		case 'p':
			cfg.monitor_pid = atoi(optarg);
			if (cfg.monitor_pid < 1) {
				fprintf(stderr, "Error: pid must be >= 1.\n");
				exit(1);
			}
			break;
		case 'o':
			cfg.output_one_time = 1;
			break;
		case 'P':
			cfg.max_processes = atoi(optarg);
			if (cfg.max_processes < 1) {
				fprintf(stderr, "Error: processes must be >= 1.\n");
				exit(1);
			}
			if (cfg.max_processes > MAX_TASKS) {
				fprintf(stderr, "Warning: processes capped to %d.\n",
					MAX_TASKS);
				cfg.max_processes = MAX_TASKS;
			}
			break;
		case 'C':
			cfg.container_path = strdup(optarg);
			break;
		case 's':
			if (strlen(optarg) == 0) {
				fprintf(stderr, "Error: empty sort field\n");
				exit(1);
			}

			field = get_field_by_name(optarg);
			/* Show available fields if invalid option provided */
			if (!field) {
				fprintf(stderr, "Error: invalid sort field '%s'\n", optarg);
				display_available_fields(MODE_TYPE_ALL);
				exit(1);
			}

			cfg.sort_field = field;
			break;
		case 'M':
			cfg.display_mode = MODE_MEMVERBOSE;
			cfg.sort_field = get_field_by_name("mem");
			break;
		default:
			fprintf(stderr, "Try 'delaytop --help' for more information.\n");
			exit(1);
		}
	}
}

/* Calculate average delay in milliseconds for overall memory */
static void set_mem_delay_total(struct task_info *t)
{
	t->mem_delay_total = t->swapin_delay_total +
		t->freepages_delay_total +
		t->thrashing_delay_total +
		t->compact_delay_total +
		t->wpcopy_delay_total;
}

static void set_mem_count(struct task_info *t)
{
	t->mem_count = t->swapin_count +
		t->freepages_count +
		t->thrashing_count +
		t->compact_count +
		t->wpcopy_count;
}

/* Create a raw netlink socket and bind */
static int create_nl_socket(void)
{
	int fd;
	struct sockaddr_nl local;

	fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
	if (fd < 0)
		return -1;

	memset(&local, 0, sizeof(local));
	local.nl_family = AF_NETLINK;

	if (bind(fd, (struct sockaddr *) &local, sizeof(local)) < 0) {
		fprintf(stderr, "Failed to bind socket when create nl_socket\n");
		close(fd);
		return -1;
	}

	return fd;
}

/* Send a command via netlink */
static int send_cmd(int sd, __u16 nlmsg_type, __u32 nlmsg_pid,
			 __u8 genl_cmd, __u16 nla_type,
			 void *nla_data, int nla_len)
{
	struct sockaddr_nl nladdr;
	struct nlattr *na;
	int r, buflen;
	char *buf;

	struct {
		struct nlmsghdr n;
		struct genlmsghdr g;
		char buf[MAX_MSG_SIZE];
	} msg;

	msg.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
	msg.n.nlmsg_type = nlmsg_type;
	msg.n.nlmsg_flags = NLM_F_REQUEST;
	msg.n.nlmsg_seq = 0;
	msg.n.nlmsg_pid = nlmsg_pid;
	msg.g.cmd = genl_cmd;
	msg.g.version = 0x1;
	na = (struct nlattr *) GENLMSG_DATA(&msg);
	na->nla_type = nla_type;
	na->nla_len = nla_len + NLA_HDRLEN;
	memcpy(NLA_DATA(na), nla_data, nla_len);
	msg.n.nlmsg_len += NLMSG_ALIGN(na->nla_len);

	buf = (char *) &msg;
	buflen = msg.n.nlmsg_len;
	memset(&nladdr, 0, sizeof(nladdr));
	nladdr.nl_family = AF_NETLINK;
	while ((r = sendto(sd, buf, buflen, 0, (struct sockaddr *) &nladdr,
					sizeof(nladdr))) < buflen) {
		if (r > 0) {
			buf += r;
			buflen -= r;
		} else if (errno != EAGAIN)
			return -1;
	}
	return 0;
}

/* Get family ID for taskstats via netlink */
static int get_family_id(int sd)
{
	struct {
		struct nlmsghdr n;
		struct genlmsghdr g;
		char buf[256];
	} ans;

	int id = 0, rc;
	struct nlattr *na;
	int rep_len;
	char name[100];

	strncpy(name, TASKSTATS_GENL_NAME, sizeof(name) - 1);
	name[sizeof(name) - 1] = '\0';
	rc = send_cmd(sd, GENL_ID_CTRL, getpid(), CTRL_CMD_GETFAMILY,
			CTRL_ATTR_FAMILY_NAME, (void *)name,
			strlen(TASKSTATS_GENL_NAME)+1);
	if (rc < 0) {
		fprintf(stderr, "Failed to send cmd for family id\n");
		return 0;
	}

	rep_len = recv(sd, &ans, sizeof(ans), 0);
	if (ans.n.nlmsg_type == NLMSG_ERROR ||
		(rep_len < 0) || !NLMSG_OK((&ans.n), rep_len)) {
		fprintf(stderr, "Failed to receive response for family id\n");
		return 0;
	}

	na = (struct nlattr *) GENLMSG_DATA(&ans);
	na = (struct nlattr *) ((char *) na + NLA_ALIGN(na->nla_len));
	if (na->nla_type == CTRL_ATTR_FAMILY_ID)
		id = *(__u16 *) NLA_DATA(na);
	return id;
}

static int read_psi_stats(void)
{
	FILE *fp;
	char line[256];
	int ret = 0;
	int error_count = 0;

	/* Check if PSI path exists */
	if (access(PSI_PATH, F_OK) != 0) {
		fprintf(stderr, "Error: PSI interface not found at %s\n", PSI_PATH);
		fprintf(stderr, "Please ensure your kernel supports PSI (Pressure Stall Information)\n");
		return -1;
	}

	/* Zero all fields */
	memset(&psi, 0, sizeof(psi));

	/* CPU pressure */
	fp = fopen(PSI_CPU_PATH, "r");
	if (fp) {
		while (fgets(line, sizeof(line), fp)) {
			if (strncmp(line, "some", 4) == 0) {
				ret = sscanf(line, "some avg10=%lf avg60=%lf avg300=%lf total=%llu",
							&psi.cpu_some_avg10, &psi.cpu_some_avg60,
							&psi.cpu_some_avg300, &psi.cpu_some_total);
				if (ret != 4) {
					fprintf(stderr, "Failed to parse CPU some PSI data\n");
					error_count++;
				}
			} else if (strncmp(line, "full", 4) == 0) {
				ret = sscanf(line, "full avg10=%lf avg60=%lf avg300=%lf total=%llu",
						&psi.cpu_full_avg10, &psi.cpu_full_avg60,
						&psi.cpu_full_avg300, &psi.cpu_full_total);
				if (ret != 4) {
					fprintf(stderr, "Failed to parse CPU full PSI data\n");
					error_count++;
				}
			}
		}
		fclose(fp);
	} else {
		fprintf(stderr, "Warning: Failed to open %s\n", PSI_CPU_PATH);
		error_count++;
	}

	/* Memory pressure */
	fp = fopen(PSI_MEMORY_PATH, "r");
	if (fp) {
		while (fgets(line, sizeof(line), fp)) {
			if (strncmp(line, "some", 4) == 0) {
				ret = sscanf(line, "some avg10=%lf avg60=%lf avg300=%lf total=%llu",
						&psi.memory_some_avg10, &psi.memory_some_avg60,
						&psi.memory_some_avg300, &psi.memory_some_total);
				if (ret != 4) {
					fprintf(stderr, "Failed to parse Memory some PSI data\n");
					error_count++;
				}
			} else if (strncmp(line, "full", 4) == 0) {
				ret = sscanf(line, "full avg10=%lf avg60=%lf avg300=%lf total=%llu",
						&psi.memory_full_avg10, &psi.memory_full_avg60,
						&psi.memory_full_avg300, &psi.memory_full_total);
				if (ret != 4) {
					fprintf(stderr, "Failed to parse Memory full PSI data\n");
					error_count++;
				}
			}
		}
		fclose(fp);
	} else {
		fprintf(stderr, "Warning: Failed to open %s\n", PSI_MEMORY_PATH);
		error_count++;
	}

	/* IO pressure */
	fp = fopen(PSI_IO_PATH, "r");
	if (fp) {
		while (fgets(line, sizeof(line), fp)) {
			if (strncmp(line, "some", 4) == 0) {
				ret = sscanf(line, "some avg10=%lf avg60=%lf avg300=%lf total=%llu",
						&psi.io_some_avg10, &psi.io_some_avg60,
						&psi.io_some_avg300, &psi.io_some_total);
				if (ret != 4) {
					fprintf(stderr, "Failed to parse IO some PSI data\n");
					error_count++;
				}
			} else if (strncmp(line, "full", 4) == 0) {
				ret = sscanf(line, "full avg10=%lf avg60=%lf avg300=%lf total=%llu",
						&psi.io_full_avg10, &psi.io_full_avg60,
						&psi.io_full_avg300, &psi.io_full_total);
				if (ret != 4) {
					fprintf(stderr, "Failed to parse IO full PSI data\n");
					error_count++;
				}
			}
		}
		fclose(fp);
	} else {
		fprintf(stderr, "Warning: Failed to open %s\n", PSI_IO_PATH);
		error_count++;
	}

	/* IRQ pressure (only full) */
	fp = fopen(PSI_IRQ_PATH, "r");
	if (fp) {
		while (fgets(line, sizeof(line), fp)) {
			if (strncmp(line, "full", 4) == 0) {
				ret = sscanf(line, "full avg10=%lf avg60=%lf avg300=%lf total=%llu",
						&psi.irq_full_avg10, &psi.irq_full_avg60,
						&psi.irq_full_avg300, &psi.irq_full_total);
				if (ret != 4) {
					fprintf(stderr, "Failed to parse IRQ full PSI data\n");
					error_count++;
				}
			}
		}
		fclose(fp);
	} else {
		fprintf(stderr, "Warning: Failed to open %s\n", PSI_IRQ_PATH);
		error_count++;
	}

	/* Return error count: 0 means success, >0 means warnings, -1 means fatal error */
	if (error_count > 0) {
		fprintf(stderr, "PSI stats reading completed with %d warnings\n", error_count);
		return error_count;
	}

	return 0;
}

static int read_comm(int pid, char *comm_buf, size_t buf_size)
{
	char path[64];
	int ret = -1;
	size_t len;
	FILE *fp;

	snprintf(path, sizeof(path), "/proc/%d/comm", pid);
	fp = fopen(path, "r");
	if (!fp) {
		fprintf(stderr, "Failed to open comm file /proc/%d/comm\n", pid);
		return ret;
	}

	if (fgets(comm_buf, buf_size, fp)) {
		len = strlen(comm_buf);
		if (len > 0 && comm_buf[len - 1] == '\n')
			comm_buf[len - 1] = '\0';
		ret = 0;
	}

	fclose(fp);

	return ret;
}

static void fetch_and_fill_task_info(int pid, const char *comm)
{
	struct {
		struct nlmsghdr n;
		struct genlmsghdr g;
		char buf[MAX_MSG_SIZE];
	} resp;
	struct taskstats stats;
	struct nlattr *nested;
	struct nlattr *na;
	int nested_len;
	int nl_len;
	int rc;

	/* Send request for task stats */
	if (send_cmd(nl_sd, family_id, getpid(), TASKSTATS_CMD_GET,
				 TASKSTATS_CMD_ATTR_PID, &pid, sizeof(pid)) < 0) {
		fprintf(stderr, "Failed to send request for task stats\n");
		return;
	}

	/* Receive response */
	rc = recv(nl_sd, &resp, sizeof(resp), 0);
	if (rc < 0 || resp.n.nlmsg_type == NLMSG_ERROR) {
		fprintf(stderr, "Failed to receive response for task stats\n");
		return;
	}

	/* Parse response */
	nl_len = GENLMSG_PAYLOAD(&resp.n);
	na = (struct nlattr *) GENLMSG_DATA(&resp);
	while (nl_len > 0) {
		if (na->nla_type == TASKSTATS_TYPE_AGGR_PID) {
			nested = (struct nlattr *) NLA_DATA(na);
			nested_len = NLA_PAYLOAD(na->nla_len);
			while (nested_len > 0) {
				if (nested->nla_type == TASKSTATS_TYPE_STATS) {
					memcpy(&stats, NLA_DATA(nested), sizeof(stats));
					if (task_count < MAX_TASKS) {
						tasks[task_count].pid = pid;
						tasks[task_count].tgid = pid;
						strncpy(tasks[task_count].command, comm,
							TASK_COMM_LEN - 1);
						tasks[task_count].command[TASK_COMM_LEN - 1] = '\0';
						SET_TASK_STAT(task_count, cpu_count);
						SET_TASK_STAT(task_count, cpu_delay_total);
						SET_TASK_STAT(task_count, blkio_count);
						SET_TASK_STAT(task_count, blkio_delay_total);
						SET_TASK_STAT(task_count, swapin_count);
						SET_TASK_STAT(task_count, swapin_delay_total);
						SET_TASK_STAT(task_count, freepages_count);
						SET_TASK_STAT(task_count, freepages_delay_total);
						SET_TASK_STAT(task_count, thrashing_count);
						SET_TASK_STAT(task_count, thrashing_delay_total);
						SET_TASK_STAT(task_count, compact_count);
						SET_TASK_STAT(task_count, compact_delay_total);
						SET_TASK_STAT(task_count, wpcopy_count);
						SET_TASK_STAT(task_count, wpcopy_delay_total);
						SET_TASK_STAT(task_count, irq_count);
						SET_TASK_STAT(task_count, irq_delay_total);
						set_mem_count(&tasks[task_count]);
						set_mem_delay_total(&tasks[task_count]);
						task_count++;
					}
					break;
				}
				nested_len -= NLA_ALIGN(nested->nla_len);
				nested = NLA_NEXT(nested);
			}
		}
		nl_len -= NLA_ALIGN(na->nla_len);
		na = NLA_NEXT(na);
	}
	return;
}

static void get_task_delays(void)
{
	char comm[TASK_COMM_LEN];
	struct dirent *entry;
	DIR *dir;
	int pid;

	task_count = 0;
	if (cfg.monitor_pid > 0) {
		if (read_comm(cfg.monitor_pid, comm, sizeof(comm)) == 0)
			fetch_and_fill_task_info(cfg.monitor_pid, comm);
		return;
	}

	dir = opendir("/proc");
	if (!dir) {
		fprintf(stderr, "Error opening /proc directory\n");
		return;
	}

	while ((entry = readdir(dir)) != NULL && task_count < MAX_TASKS) {
		if (!isdigit(entry->d_name[0]))
			continue;
		pid = atoi(entry->d_name);
		if (pid == 0)
			continue;
		if (read_comm(pid, comm, sizeof(comm)) != 0)
			continue;
		fetch_and_fill_task_info(pid, comm);
	}
	closedir(dir);
}

/* Calculate average delay in milliseconds */
static double average_ms(unsigned long long total, unsigned long long count)
{
	if (count == 0)
		return 0;
	return (double)total / 1000000.0 / count;
}

/* Comparison function for sorting tasks */
static int compare_tasks(const void *a, const void *b)
{
	const struct task_info *t1 = (const struct task_info *)a;
	const struct task_info *t2 = (const struct task_info *)b;
	unsigned long long total1;
	unsigned long long total2;
	unsigned long count1;
	unsigned long count2;
	double avg1, avg2;

	total1 = *(unsigned long long *)((char *)t1 + cfg.sort_field->total_offset);
	total2 = *(unsigned long long *)((char *)t2 + cfg.sort_field->total_offset);
	count1 = *(unsigned long *)((char *)t1 + cfg.sort_field->count_offset);
	count2 = *(unsigned long *)((char *)t2 + cfg.sort_field->count_offset);

	avg1 = average_ms(total1, count1);
	avg2 = average_ms(total2, count2);
	if (avg1 != avg2)
		return avg2 > avg1 ? 1 : -1;

	return 0;
}

/* Sort tasks by selected field */
static void sort_tasks(void)
{
	if (task_count > 0)
		qsort(tasks, task_count, sizeof(struct task_info), compare_tasks);
}

/* Get container statistics via cgroupstats */
static void get_container_stats(void)
{
	int rc, cfd;
	struct {
		struct nlmsghdr n;
		struct genlmsghdr g;
		char buf[MAX_MSG_SIZE];
	} req, resp;
	struct nlattr *na;
	int nl_len;
	struct cgroupstats stats;

	/* Check if container path is set */
	if (!cfg.container_path)
		return;

	/* Open container cgroup */
	cfd = open(cfg.container_path, O_RDONLY);
	if (cfd < 0) {
		fprintf(stderr, "Error opening container path: %s\n", cfg.container_path);
		return;
	}

	/* Send request for container stats */
	if (send_cmd(nl_sd, family_id, getpid(), CGROUPSTATS_CMD_GET,
				CGROUPSTATS_CMD_ATTR_FD, &cfd, sizeof(__u32)) < 0) {
		fprintf(stderr, "Failed to send request for container stats\n");
		close(cfd);
		return;
	}

	/* Receive response */
	rc = recv(nl_sd, &resp, sizeof(resp), 0);
	if (rc < 0 || resp.n.nlmsg_type == NLMSG_ERROR) {
		fprintf(stderr, "Failed to receive response for container stats\n");
		close(cfd);
		return;
	}

	/* Parse response */
	nl_len = GENLMSG_PAYLOAD(&resp.n);
	na = (struct nlattr *) GENLMSG_DATA(&resp);
	while (nl_len > 0) {
		if (na->nla_type == CGROUPSTATS_TYPE_CGROUP_STATS) {
			/* Get the cgroupstats structure */
			memcpy(&stats, NLA_DATA(na), sizeof(stats));

			/* Fill container stats */
			container_stats.nr_sleeping = stats.nr_sleeping;
			container_stats.nr_running = stats.nr_running;
			container_stats.nr_stopped = stats.nr_stopped;
			container_stats.nr_uninterruptible = stats.nr_uninterruptible;
			container_stats.nr_io_wait = stats.nr_io_wait;
			break;
		}
		nl_len -= NLA_ALIGN(na->nla_len);
		na = (struct nlattr *) ((char *) na + NLA_ALIGN(na->nla_len));
	}

	close(cfd);
}

/* Display results to stdout or log file */
static void display_results(int psi_ret)
{
	time_t now = time(NULL);
	struct tm *tm_now = localtime(&now);
	FILE *out = stdout;
	char timestamp[32];
	bool suc = true;
	int i, count;

	/* Clear terminal screen */
	suc &= BOOL_FPRINT(out, "\033[H\033[J");

	/* PSI output (one-line, no cat style) */
	suc &= BOOL_FPRINT(out, "System Pressure Information: (avg10/avg60vg300/total)\n");
	if (psi_ret) {
		suc &= BOOL_FPRINT(out, "  PSI not found: check if psi=1 enabled in cmdline\n");
	} else {
		suc &= BOOL_FPRINT(out, PSI_LINE_FORMAT,
			"CPU some:",
			psi.cpu_some_avg10,
			psi.cpu_some_avg60,
			psi.cpu_some_avg300,
			psi.cpu_some_total / 1000);
		suc &= BOOL_FPRINT(out, PSI_LINE_FORMAT,
			"CPU full:",
			psi.cpu_full_avg10,
			psi.cpu_full_avg60,
			psi.cpu_full_avg300,
			psi.cpu_full_total / 1000);
		suc &= BOOL_FPRINT(out, PSI_LINE_FORMAT,
			"Memory full:",
			psi.memory_full_avg10,
			psi.memory_full_avg60,
			psi.memory_full_avg300,
			psi.memory_full_total / 1000);
		suc &= BOOL_FPRINT(out, PSI_LINE_FORMAT,
			"Memory some:",
			psi.memory_some_avg10,
			psi.memory_some_avg60,
			psi.memory_some_avg300,
			psi.memory_some_total / 1000);
		suc &= BOOL_FPRINT(out, PSI_LINE_FORMAT,
			"IO full:",
			psi.io_full_avg10,
			psi.io_full_avg60,
			psi.io_full_avg300,
			psi.io_full_total / 1000);
		suc &= BOOL_FPRINT(out, PSI_LINE_FORMAT,
			"IO some:",
			psi.io_some_avg10,
			psi.io_some_avg60,
			psi.io_some_avg300,
			psi.io_some_total / 1000);
		suc &= BOOL_FPRINT(out, PSI_LINE_FORMAT,
			"IRQ full:",
			psi.irq_full_avg10,
			psi.irq_full_avg60,
			psi.irq_full_avg300,
			psi.irq_full_total / 1000);
	}

	if (cfg.container_path) {
		suc &= BOOL_FPRINT(out, "Container Information (%s):\n", cfg.container_path);
		suc &= BOOL_FPRINT(out, "Processes: running=%d, sleeping=%d, ",
			container_stats.nr_running, container_stats.nr_sleeping);
		suc &= BOOL_FPRINT(out, "stopped=%d, uninterruptible=%d, io_wait=%d\n\n",
			container_stats.nr_stopped, container_stats.nr_uninterruptible,
			container_stats.nr_io_wait);
	}

	/* Interacive command */
	suc &= BOOL_FPRINT(out, "[o]sort [M]memverbose [q]quit\n");
	if (sort_selected) {
		if (cfg.display_mode == MODE_MEMVERBOSE)
			suc &= BOOL_FPRINT(out,
				"sort selection: [m]MEM [r]RCL [t]THR [p]CMP [w]WP\n");
		else
			suc &= BOOL_FPRINT(out,
				"sort selection: [c]CPU [i]IO [m]MEM [q]IRQ\n");
	}

	/* Task delay output */
	suc &= BOOL_FPRINT(out, "Top %d processes (sorted by %s delay):\n",
			cfg.max_processes, get_name_by_field(cfg.sort_field));

	suc &= BOOL_FPRINT(out, "%8s  %8s  %-17s", "PID", "TGID", "COMMAND");
	if (cfg.display_mode == MODE_MEMVERBOSE) {
		suc &= BOOL_FPRINT(out, "%8s %8s %8s %8s %8s %8s\n",
			"MEM(ms)", "SWAP(ms)", "RCL(ms)",
			"THR(ms)", "CMP(ms)", "WP(ms)");
		suc &= BOOL_FPRINT(out, "-----------------------");
		suc &= BOOL_FPRINT(out, "-----------------------");
		suc &= BOOL_FPRINT(out, "-----------------------");
		suc &= BOOL_FPRINT(out, "---------------------\n");
	} else {
		suc &= BOOL_FPRINT(out, "%8s %8s %8s %8s\n",
			"CPU(ms)", "IO(ms)", "IRQ(ms)", "MEM(ms)");
		suc &= BOOL_FPRINT(out, "-----------------------");
		suc &= BOOL_FPRINT(out, "-----------------------");
		suc &= BOOL_FPRINT(out, "--------------------------\n");
	}

	count = task_count < cfg.max_processes ? task_count : cfg.max_processes;

	for (i = 0; i < count; i++) {
		suc &= BOOL_FPRINT(out, "%8d  %8d  %-15s",
			tasks[i].pid, tasks[i].tgid, tasks[i].command);
		if (cfg.display_mode == MODE_MEMVERBOSE) {
			suc &= BOOL_FPRINT(out, DELAY_FMT_MEMVERBOSE,
				TASK_AVG(tasks[i], mem),
				TASK_AVG(tasks[i], swapin),
				TASK_AVG(tasks[i], freepages),
				TASK_AVG(tasks[i], thrashing),
				TASK_AVG(tasks[i], compact),
				TASK_AVG(tasks[i], wpcopy));
		} else {
			suc &= BOOL_FPRINT(out, DELAY_FMT_DEFAULT,
				TASK_AVG(tasks[i], cpu),
				TASK_AVG(tasks[i], blkio),
				TASK_AVG(tasks[i], irq),
				TASK_AVG(tasks[i], mem));
		}
	}

	suc &= BOOL_FPRINT(out, "\n");

	if (!suc)
		perror("Error writing to output");
}

/* Check for keyboard input with timeout based on cfg.delay */
static char check_for_keypress(void)
{
	struct timeval tv = {cfg.delay, 0};
	fd_set readfds;
	char ch = 0;

	FD_ZERO(&readfds);
	FD_SET(STDIN_FILENO, &readfds);
	int r = select(STDIN_FILENO + 1, &readfds, NULL, NULL, &tv);

	if (r > 0 && FD_ISSET(STDIN_FILENO, &readfds)) {
		read(STDIN_FILENO, &ch, 1);
		return ch;
	}

	return 0;
}

#define MAX_MODE_SIZE 2
static void toggle_display_mode(void)
{
	static const size_t modes[MAX_MODE_SIZE] = {MODE_DEFAULT, MODE_MEMVERBOSE};
	static size_t cur_index;

	cur_index = (cur_index + 1) % MAX_MODE_SIZE;
	cfg.display_mode = modes[cur_index];
}

/* Handle keyboard input: sorting selection, mode toggle, or quit */
static void handle_keypress(char ch, int *running)
{
	const struct field_desc *field;

	/* Change sort field */
	if (sort_selected) {
		field = get_field_by_cmd_char(ch);
		if (field && (field->supported_modes & cfg.display_mode))
			cfg.sort_field = field;

		sort_selected = 0;
	/* Handle mode changes or quit */
	} else {
		switch (ch) {
		case 'o':
			sort_selected = 1;
			break;
		case 'M':
			toggle_display_mode();
			for (field = sort_fields; field->name != NULL; field++) {
				if (field->supported_modes & cfg.display_mode) {
					cfg.sort_field = field;
					break;
				}
			}
			break;
		case 'q':
		case 'Q':
			*running = 0;
			break;
		default:
			break;
		}
	}
}

/* Main function */
int main(int argc, char **argv)
{
	const struct field_desc *field;
	int iterations = 0;
	int psi_ret = 0;
	char keypress;

	/* Parse command line arguments */
	parse_args(argc, argv);

	/* Setup netlink socket */
	nl_sd = create_nl_socket();
	if (nl_sd < 0) {
		fprintf(stderr, "Error creating netlink socket\n");
		exit(1);
	}

	/* Get family ID for taskstats via netlink */
	family_id = get_family_id(nl_sd);
	if (!family_id) {
		fprintf(stderr, "Error getting taskstats family ID\n");
		close(nl_sd);
		exit(1);
	}

	/* Set terminal to non-canonical mode for interaction */
	enable_raw_mode();

	/* Main loop */
	while (running) {
		/* Auto-switch sort field when not matching display mode */
		if (!(cfg.sort_field->supported_modes & cfg.display_mode)) {
			for (field = sort_fields; field->name != NULL; field++) {
				if (field->supported_modes & cfg.display_mode) {
					cfg.sort_field = field;
					printf("Auto-switched sort field to: %s\n", field->name);
					break;
				}
			}
		}

		/* Read PSI statistics */
		psi_ret = read_psi_stats();

		/* Get container stats if container path provided */
		if (cfg.container_path)
			get_container_stats();

		/* Get task delays */
		get_task_delays();

		/* Sort tasks */
		sort_tasks();

		/* Display results to stdout or log file */
		display_results(psi_ret);

		/* Check for iterations */
		if (cfg.iterations > 0 && ++iterations >= cfg.iterations)
			break;

		/* Exit if output_one_time is set */
		if (cfg.output_one_time)
			break;

		/* Keypress for interactive usage */
		keypress = check_for_keypress();
		if (keypress)
			handle_keypress(keypress, &running);
	}

	/* Restore terminal mode */
	disable_raw_mode();

	/* Cleanup */
	close(nl_sd);
	if (cfg.container_path)
		free(cfg.container_path);

	return 0;
}